[dace] Add offset to scan range

GridTools · Oct 11, 2023 · 23f4d5f · 23f4d5f
1 parent 397c376
commit 23f4d5f
Showing 1 changed file with 3 additions and 15 deletions.
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py
@@ -290,21 +290,15 @@ def visit_StencilClosure(
             _, (scan_lb, scan_ub) = closure_domain[scan_dim_index]
             output_subset = f"{scan_lb.value}:{scan_ub.value}"
 
-            closure_sdfg.add_array(
-                nsdfg_output_name,
-                dtype=output_descriptor.dtype,
-                shape=(array_table[output_name].shape[scan_dim_index],),
-                strides=(array_table[output_name].strides[scan_dim_index],),
-                transient=True,
-            )
-
+            scan_offset = closure_sdfg.arrays[output_name].offset[scan_dim_index]
+            scan_shape = closure_sdfg.arrays[output_name].shape[scan_dim_index]
             output_memlet = create_memlet_at(
                 output_name,
                 self.storage_types[output_name],
                 {
                     dim: f"i_{dim}"
                     if f"i_{dim}" in map_ranges
-                    else f"0:{output_descriptor.shape[scan_dim_index]}"
+                    else f"{scan_offset}:{scan_offset}+{scan_shape}"
                     for dim, _ in closure_domain
                 },
             )
@@ -316,12 +310,6 @@ def visit_StencilClosure(
 
             output_subset = "0"
 
-            closure_sdfg.add_scalar(
-                nsdfg_output_name,
-                dtype=output_descriptor.dtype,
-                transient=True,
-            )
-
             output_memlet = create_memlet_at(
                 output_name,
                 self.storage_types[output_name],