From 23f4d5fcfd202b49570ed5fb8a9b4d76f44f2b48 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 11 Oct 2023 15:10:13 +0200 Subject: [PATCH] [dace] Add offset to scan range --- .../runners/dace_iterator/itir_to_sdfg.py | 18 +++--------------- 1 file changed, 3 insertions(+), 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 index 78350eb94b..7025ccf002 100644 --- 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],