Repository: GridTools/gt4py
Pull request #1668 (2024-09-27 12:54:59)

fix[next][dace]: enable subdomain in let-statements
This PR fixes the lowering of let-statements with fields defined on sub-domain. The previous implementation was allocating dace array with the exact domain of the computation, eventually not starting from index 0. This required the use of offsets in dace arrays to compensate for the start index. There were two problems with this: 1) in the lowering from GTIR to SDFG because the array offsets were not correctly propagated across nested SDFGs for let-statements; and 2) the use of array offsets caused issue to SDFG-inlining pass. In order to reduce the overall complexity, this PR proposes to remove the usage of array offsets and therefore always allocate temporary fields starting from index 0.


Patch:
diff --git a/src/gt4py/next/program_processors/runners/dace_common/workflow.py b/src/gt4py/next/program_processors/runners/dace_common/workflow.py
index 1caa4684..f7809e3c 100644
--- a/src/gt4py/next/program_processors/runners/dace_common/workflow.py
+++ b/src/gt4py/next/program_processors/runners/dace_common/workflow.py
@@ -121,7 +121,7 @@ def convert_args(
             # the arguments list has to be reconstructed.
             for i, (arg_name, arg_type) in enumerate(inp.sdfg_arglist):
                 if isinstance(arg_type, dace.data.Array):
-                    assert arg_name in kwargs, f"Argument '{arg_name}' not found."
+                    assert arg_name in kwargs, f"argument '{arg_name}' not found."
                     data_ptr = get_array_interface_ptr(kwargs[arg_name], arg_type.storage)
                     assert isinstance(last_call_args[i], ctypes.c_void_p)
                     if last_call_args[i].value != data_ptr:
@@ -138,7 +138,7 @@ def convert_args(
                         # shape and strides of arrays are supposed not to change, and can therefore be omitted
                         assert dace_utils.is_field_symbol(
                             arg_name
-                        ), f"Argument '{arg_name}' not found."
+                        ), f"argument '{arg_name}' not found."
 
             if use_fast_call:
                 return sdfg_program.fast_call(*sdfg_program._lastargs)
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
index 28085107..5c835095 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
@@ -117,16 +117,17 @@ def _create_temporary_field(
     node_type: ts.FieldType,
     output_desc: dace.data.Data,
 ) -> tuple[dace.nodes.AccessNode, ts.FieldType]:
-    domain_dims, domain_lbs, domain_ubs = zip(*domain)
+    domain_dims, _, domain_ubs = zip(*domain)
     field_dims = list(domain_dims)
-    field_shape = [
-        # diff between upper and lower bound
-        (ub - lb)
-        for lb, ub in zip(domain_lbs, domain_ubs)
-    ]
-    field_offset: Optional[list[dace.symbolic.SymbolicType]] = None
-    if any(domain_lbs):
-        field_offset = [-lb for lb in domain_lbs]
+    # It should be enough to allocate an array with shape (upper_bound - lower_bound)
+    # but this would require to use array offset for compensate for the start index.
+    # Suppose that a field operator executes on domain [2,N-2], the dace array to store
+    # the result only needs size (N-4), but this would require to compensate all array
+    # accesses with offset -2 (which corresponds to -lower_bound). Instead, we choose
+    # to allocate (N-2), leaving positions [0:2] unused. The reason is that array offset
+    # is known to cause issues to SDFG inlining. Besides, map fusion will in any case
+    # eliminate most of transient arrays.
+    field_shape = list(domain_ubs)
 
     if isinstance(output_desc, dace.data.Array):
         assert isinstance(node_type.dtype, gtir_ts.ListType)
@@ -139,9 +140,7 @@ def _create_temporary_field(
 
     # allocate local temporary storage for the result field
     assert isinstance(field_dtype, ts.ScalarType)
-    temp_name, _ = sdfg.add_temp_transient(
-        field_shape, dace_utils.as_dace_type(field_dtype), offset=field_offset
-    )
+    temp_name, _ = sdfg.add_temp_transient(field_shape, dace_utils.as_dace_type(field_dtype))
     field_node = state.add_access(temp_name)
     field_type = ts.FieldType(field_dims, node_type.dtype)
 
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_to_sdfg.py
index f702a633..8617dc7c 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_to_sdfg.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_to_sdfg.py
@@ -25,6 +25,7 @@ from gt4py.eve import concepts
 from gt4py.next import common as gtx_common
 from gt4py.next.iterator import ir as gtir
 from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm
+from gt4py.next.iterator.transforms import symbol_ref_utils
 from gt4py.next.iterator.type_system import inference as gtir_type_inference
 from gt4py.next.program_processors.runners.dace_common import utility as dace_utils
 from gt4py.next.program_processors.runners.dace_fieldview import (
@@ -400,9 +401,11 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         lambda_args_mapping = {str(p.id): arg for p, arg in zip(node.params, args, strict=True)}
 
         # inherit symbols from parent scope but eventually override with local symbols
-        lambda_symbols = self.global_symbols | {
-            pname: type_ for pname, (_, type_) in lambda_args_mapping.items()
+        lambda_symbols = {
+            sym: self.global_symbols[sym]
+            for sym in symbol_ref_utils.collect_symbol_refs(node.expr, self.global_symbols.keys())
         }
+        lambda_symbols |= {pname: type_ for pname, (_, type_) in lambda_args_mapping.items()}
 
         nsdfg = dace.SDFG(f"{sdfg.label}_nested")
         nstate = nsdfg.add_state("lambda")
@@ -436,25 +439,25 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
                 src_node, _ = lambda_args_mapping[nsdfg_dataname]
                 dataname = src_node.data
                 datadesc = src_node.desc(sdfg)
-
-                nsdfg_symbols_mapping |= {
-                    str(nested_symbol): parent_symbol
-                    for nested_symbol, parent_symbol in zip(
-                        [*nsdfg_datadesc.shape, *nsdfg_datadesc.strides],
-                        [*datadesc.shape, *datadesc.strides],
-                        strict=True,
-                    )
-                    if isinstance(nested_symbol, dace.symbol)
-                }
             else:
                 dataname = nsdfg_dataname
                 datadesc = sdfg.arrays[nsdfg_dataname]
-                # ensure that connectivity tables are non-transient arrays in parent SDFG
-                if dataname in connectivity_arrays:
-                    datadesc.transient = False
 
-            if datadesc:
-                input_memlets[nsdfg_dataname] = dace.Memlet.from_array(dataname, datadesc)
+            # ensure that connectivity tables are non-transient arrays in parent SDFG
+            if dataname in connectivity_arrays:
+                datadesc.transient = False
+
+            input_memlets[nsdfg_dataname] = sdfg.make_array_memlet(dataname)
+
+            nsdfg_symbols_mapping |= {
+                str(nested_symbol): parent_symbol
+                for nested_symbol, parent_symbol in zip(
+                    [*nsdfg_datadesc.shape, *nsdfg_datadesc.strides],
+                    [*datadesc.shape, *datadesc.strides],
+                    strict=True,
+                )
+                if isinstance(nested_symbol, dace.symbol)
+            }
 
         nsdfg_node = head_state.add_nested_sdfg(
             nsdfg,
@@ -484,9 +487,7 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
                 nstate.remove_node(lambda_node)
             temp, _ = sdfg.add_temp_transient_like(desc)
             dst_node = head_state.add_access(temp)
-            head_state.add_edge(
-                nsdfg_node, connector, dst_node, None, dace.Memlet.from_array(temp, desc)
-            )
+            head_state.add_edge(nsdfg_node, connector, dst_node, None, sdfg.make_array_memlet(temp))
             results.append((dst_node, type_))
 
         return results


Test patch:
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
index 694cf731..0c446d97 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
@@ -1231,6 +1231,9 @@ def test_gtir_let_lambda():
     domain = im.call("cartesian_domain")(
         im.call("named_range")(gtir.AxisLiteral(value=IDim.value), 0, "size")
     )
+    subdomain = im.call("cartesian_domain")(
+        im.call("named_range")(gtir.AxisLiteral(value=IDim.value), 1, im.minus("size", 1))
+    )
     testee = gtir.Program(
         id="let_lambda",
         function_definitions=[],
@@ -1246,21 +1249,21 @@ def test_gtir_let_lambda():
                 # `x2` is a let-lambda expression representing `x * 4`
                 #  - note that the let-symbol `x2` is used twice, in a nested let-expression, to test aliasing of the symbol
                 # `x3` is a let-lambda expression simply accessing `x` field symref
-                expr=im.let("x1", im.op_as_fieldop("multiplies", domain)(3.0, "x"))(
+                expr=im.let("x1", im.op_as_fieldop("multiplies", subdomain)(3.0, "x"))(
                     im.let(
                         "x2",
                         im.let("x2", im.op_as_fieldop("multiplies", domain)(2.0, "x"))(
-                            im.op_as_fieldop("plus", domain)("x2", "x2")
+                            im.op_as_fieldop("plus", subdomain)("x2", "x2")
                         ),
                     )(
                         im.let("x3", "x")(
-                            im.op_as_fieldop("plus", domain)(
-                                "x1", im.op_as_fieldop("plus", domain)("x2", "x3")
+                            im.op_as_fieldop("plus", subdomain)(
+                                "x1", im.op_as_fieldop("plus", subdomain)("x2", "x3")
                             )
                         )
                     )
                 ),
-                domain=domain,
+                domain=subdomain,
                 target=gtir.SymRef(id="y"),
             )
         ],
@@ -1268,11 +1271,12 @@ def test_gtir_let_lambda():
 
     a = np.random.rand(N)
     b = np.empty_like(a)
+    ref = np.concatenate((b[0:1], a[1 : N - 1] * 8, b[N - 1 : N]))
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, {})
 
     sdfg(a, b, **FSYMBOLS)
-    assert np.allclose(b, a * 8)
+    assert np.allclose(b, ref)
 
 
 def test_gtir_let_lambda_with_connectivity():