From a93e8dc29639265dd4ddaf3fd1c60d7d5592f67c Mon Sep 17 00:00:00 2001 From: Enrique Gonzalez Paredes Date: Wed, 25 Oct 2023 15:44:30 +0200 Subject: [PATCH] More fixes --- src/gt4py/next/program_processors/runners/roundtrip.py | 5 ++++- src/gt4py/storage/cartesian/utils.py | 4 ++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/program_processors/runners/roundtrip.py b/src/gt4py/next/program_processors/runners/roundtrip.py index 45bb47e73d..065a32a14d 100644 --- a/src/gt4py/next/program_processors/runners/roundtrip.py +++ b/src/gt4py/next/program_processors/runners/roundtrip.py @@ -216,9 +216,12 @@ def execute_roundtrip( return fencil(*args, **new_kwargs) -class RoundtripExecutor(ppi.ProgramExecutor, next_allocators.DefaultCPUAllocator): +class RoundtripExecutor(ppi.ProgramExecutor): def __call__(self, program: itir.FencilDefinition, *args, **kwargs) -> None: execute_roundtrip(program, *args, **kwargs) + __gt_device_type__ = next_allocators.DefaultCPUAllocator.__gt_device_type__ + __gt_allocate__ = next_allocators.DefaultCPUAllocator.__gt_allocate__ + executor = RoundtripExecutor() diff --git a/src/gt4py/storage/cartesian/utils.py b/src/gt4py/storage/cartesian/utils.py index de4121d2c0..1884835717 100644 --- a/src/gt4py/storage/cartesian/utils.py +++ b/src/gt4py/storage/cartesian/utils.py @@ -220,7 +220,7 @@ def allocate_cpu( buffer = _CPUBufferAllocator.allocate( shape, core_defs.dtype(dtype), - device_id=device.id, + device_id=device.device_id, layout_map=layout_map, byte_alignment=alignment_bytes, aligned_index=aligned_index, @@ -235,7 +235,7 @@ def allocate_gpu( alignment_bytes: int, aligned_index: Optional[Sequence[int]], ) -> Tuple["cp.ndarray", "cp.ndarray"]: - assert _GPUBufferAllocator is None, "GPU allocation library or device not found" + assert _GPUBufferAllocator is not None, "GPU allocation library or device not found" device = core_defs.Device( # type: ignore[type-var] core_defs.DeviceType.ROCM if gt_config.GT4PY_USE_HIP else core_defs.DeviceType.CUDA, 0 )