Fix gtfn gpu backend and executor definitions

GridTools · Nov 14, 2023 · 7756a8e · 7756a8e
1 parent bf9caa0
commit 7756a8e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/src/gt4py/next/allocators.py b/src/gt4py/next/allocators.py
@@ -168,7 +168,7 @@ def __gt_allocate__(
         shape = domain.shape
         layout_map = self.layout_mapper(domain.dims)
         # TODO(egparedes): add support for non-empty aligned index values
-        assert aligned_index is None  
+        assert aligned_index is None
 
         return core_allocators.NDArrayBufferAllocator(self.device_type, self.array_ns).allocate(
             shape, dtype, device_id, layout_map, self.byte_alignment, aligned_index

diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py
@@ -174,8 +174,22 @@ def compilation_hash(otf_closure: stages.ProgramCall) -> int:
     allocator=next_allocators.StandardCPUFieldBufferAllocator(),
 )
 
-
+gtfn_gpu_executor = otf_compile_executor.OTFCompileExecutor(
+    name="run_gtfn_gpu", otf_workflow=GTFN_GPU_WORKFLOW
+)
 run_gtfn_gpu = otf_compile_executor.OTFBackend(
-    executor=gtfn_cached_executor,
+    executor=gtfn_gpu_executor,
     allocator=next_allocators.StandardGPUFieldBufferAllocator(),
 )
+
+
+gtfn_gpu_cached_executor = otf_compile_executor.CachedOTFCompileExecutor(
+    name="run_gtfn_gpu_cached",
+    otf_workflow=workflow.CachedStep(
+        step=gtfn_gpu_executor.otf_workflow, hash_function=compilation_hash
+    ),
+)
+run_gtfn_gpu_cached = otf_compile_executor.OTFBackend(
+    executor=gtfn_gpu_cached_executor,
+    allocator=next_allocators.StandardCPUFieldBufferAllocator(),
+)