Repository: GridTools/gt4py
Pull request #1609 (2024-08-08 20:43:29)

feat[next][dace]: Fix GTIR-to-SDFG lowering for failing tests
This PR addresses some problems encountered during GTIR/DaCe integration.
- Handle field size symbol as explicit scalar argument
- Support SDFG positional arguments (not only keyword arguments)
- Use explicit subset on source node for target write
- Add test case for copy from one field to itself using an explicit domain subset
- Extend test case with `cond` by using an expression on scalar arguments as if-statement


Patch:
diff --git a/src/gt4py/next/iterator/transforms/constant_folding.py b/src/gt4py/next/iterator/transforms/constant_folding.py
index 8fdbffa9..2084ab25 100644
--- a/src/gt4py/next/iterator/transforms/constant_folding.py
+++ b/src/gt4py/next/iterator/transforms/constant_folding.py
@@ -43,9 +43,15 @@ class ConstantFolding(PreserveLocationVisitor, NodeTranslator):
             and len(new_node.args) > 0
             and all(isinstance(arg, ir.Literal) for arg in new_node.args)
         ):  # `1 + 1` -> `2`
-            if new_node.fun.id in ir.ARITHMETIC_BUILTINS:
-                fun = getattr(embedded, str(new_node.fun.id))
-                arg_values = [getattr(embedded, str(arg.type))(arg.value) for arg in new_node.args]  # type: ignore[attr-defined] # arg type already established in if condition
-                new_node = im.literal_from_value(fun(*arg_values))
+            try:
+                if new_node.fun.id in ir.ARITHMETIC_BUILTINS:
+                    fun = getattr(embedded, str(new_node.fun.id))
+                    arg_values = [
+                        getattr(embedded, str(arg.type))(arg.value)  # type: ignore[attr-defined] # arg type already established in if condition
+                        for arg in new_node.args
+                    ]
+                    new_node = im.literal_from_value(fun(*arg_values))
+            except ValueError:
+                pass  # happens for inf and neginf
 
         return new_node
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 a3834aad..bf1a0173 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
@@ -143,12 +143,14 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
             (
                 neighbor_tables[dim.value].max_neighbors
                 if dim.kind == gtx_common.DimensionKind.LOCAL
-                # we reuse the same symbol for field size passed as scalar argument to the gt4py program
-                else dace.symbol(f"__{name}_size_{i}", dtype)
+                else dace.symbol(dace_fieldview_util.field_size_symbol_name(name, i), dtype)
             )
             for i, dim in enumerate(dims)
         ]
-        strides = [dace.symbol(f"__{name}_stride_{i}", dtype) for i in range(len(dims))]
+        strides = [
+            dace.symbol(dace_fieldview_util.field_stride_symbol_name(name, i), dtype)
+            for i in range(len(dims))
+        ]
         return shape, strides
 
     def _add_storage(
@@ -170,8 +172,14 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
             sdfg.add_array(name, sym_shape, dtype, strides=sym_strides, transient=transient)
         elif isinstance(symbol_type, ts.ScalarType):
             dtype = dace_fieldview_util.as_dace_type(symbol_type)
-            # scalar arguments passed to the program are represented as symbols in DaCe SDFG
-            sdfg.add_symbol(name, dtype)
+            # Scalar arguments passed to the program are represented as symbols in DaCe SDFG.
+            # The field size is sometimes passed as scalar argument to the program, so we have to
+            # check if the shape symbol was already allocated by `_make_array_shape_and_strides`.
+            # We assume that the scalar argument for field size always follows the field argument.
+            if name in sdfg.symbols:
+                assert sdfg.symbols[name].dtype == dtype
+            else:
+                sdfg.add_symbol(name, dtype)
         else:
             raise RuntimeError(f"Data type '{type(symbol_type)}' not supported.")
 
@@ -274,6 +282,12 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
             head_state._debuginfo = dace_fieldview_util.debug_info(stmt, default=sdfg.debuginfo)
             self.visit(stmt, sdfg=sdfg, state=head_state)
 
+        # Create the call signature for the SDFG.
+        #  Only the arguments required by the GT4Py program, i.e. `node.params`, are added
+        #  as positional arguments. The implicit arguments, such as the offset providers or
+        #  the arguments created by the translation process, must be passed as keyword arguments.
+        sdfg.arg_names = [str(a) for a in node.params]
+
         sdfg.validate()
         return sdfg
 
@@ -309,7 +323,7 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
             state.add_nedge(
                 expr_node,
                 target_node,
-                dace.Memlet(data=target_node.data, subset=subset),
+                dace.Memlet(data=target_node.data, subset=subset, other_subset=subset),
             )
 
     def visit_FunCall(
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py b/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
index 8a963e80..43541925 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
@@ -73,6 +73,14 @@ def debug_info(
     return default
 
 
+def field_size_symbol_name(field_name: str, axis: int) -> str:
+    return f"__{field_name}_size_{axis}"
+
+
+def field_stride_symbol_name(field_name: str, axis: int) -> str:
+    return f"__{field_name}_stride_{axis}"
+
+
 def filter_connectivities(offset_provider: Mapping[str, Any]) -> dict[str, gtx_common.Connectivity]:
     """
     Filter offset providers of type `Connectivity`.


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
index 83a9b802..0bf8dcb6 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
@@ -46,3 +46,17 @@ def test_constant_folding_minimum():
     expected = im.ref("a")
     actual = ConstantFolding.apply(testee)
     assert actual == expected
+
+
+def test_constant_folding_literal():
+    testee = im.plus(im.literal_from_value(1), im.literal_from_value(2))
+    expected = im.literal_from_value(3)
+    actual = ConstantFolding.apply(testee)
+    assert actual == expected
+
+
+def test_constant_folding_literal_maximum():
+    testee = im.call("maximum")(im.literal_from_value(1), im.literal_from_value(2))
+    expected = im.literal_from_value(2)
+    actual = ConstantFolding.apply(testee)
+    assert actual == expected
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_dace_fieldview.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_dace_fieldview.py
index a8fc9eb3..55b037b1 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_dace_fieldview.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_dace_fieldview.py
@@ -133,10 +133,40 @@ def test_gtir_cast():
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
 
-    sdfg(x=a, y=b, z=c, **FSYMBOLS)
+    sdfg(a, b, c, **FSYMBOLS)
     np.testing.assert_array_equal(c, True)
 
 
+def test_gtir_copy_self():
+    domain = im.call("cartesian_domain")(
+        im.call("named_range")(gtir.AxisLiteral(value=IDim.value), 1, 2)
+    )
+    testee = gtir.Program(
+        id="gtir_copy_self",
+        function_definitions=[],
+        params=[
+            gtir.Sym(id="x", type=IFTYPE),
+            gtir.Sym(id="size", type=SIZE_TYPE),
+        ],
+        declarations=[],
+        body=[
+            gtir.SetAt(
+                expr=gtir.SymRef(id="x"),
+                domain=domain,
+                target=gtir.SymRef(id="x"),
+            )
+        ],
+    )
+
+    a = np.random.rand(N)
+    ref = a.copy()
+
+    sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
+
+    sdfg(a, **FSYMBOLS)
+    assert np.allclose(a, ref)
+
+
 def test_gtir_update():
     domain = im.call("cartesian_domain")(
         im.call("named_range")(gtir.AxisLiteral(value=IDim.value), 0, "size")
@@ -169,7 +199,7 @@ def test_gtir_update():
         a = np.random.rand(N)
         ref = a + 1.0
 
-        sdfg(x=a, **FSYMBOLS)
+        sdfg(a, **FSYMBOLS)
         assert np.allclose(a, ref)
 
 
@@ -202,7 +232,7 @@ def test_gtir_sum2():
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
 
-    sdfg(x=a, y=b, z=c, **FSYMBOLS)
+    sdfg(a, b, c, **FSYMBOLS)
     assert np.allclose(c, (a + b))
 
 
@@ -233,7 +263,7 @@ def test_gtir_sum2_sym():
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, {})
 
-    sdfg(x=a, z=b, **FSYMBOLS)
+    sdfg(a, b, **FSYMBOLS)
     assert np.allclose(b, (a + a))
 
 
@@ -279,7 +309,7 @@ def test_gtir_sum3():
 
         d = np.empty_like(a)
 
-        sdfg(x=a, y=b, w=c, z=d, **FSYMBOLS)
+        sdfg(a, b, c, d, **FSYMBOLS)
         assert np.allclose(d, (a + b + c))
 
 
@@ -295,7 +325,8 @@ def test_gtir_cond():
             gtir.Sym(id="y", type=IFTYPE),
             gtir.Sym(id="w", type=IFTYPE),
             gtir.Sym(id="z", type=IFTYPE),
-            gtir.Sym(id="pred", type=ts.ScalarType(ts.ScalarKind.BOOL)),
+            gtir.Sym(id="s1", type=ts.ScalarType(ts.ScalarKind.INT32)),
+            gtir.Sym(id="s2", type=ts.ScalarType(ts.ScalarKind.INT32)),
             gtir.Sym(id="scalar", type=ts.ScalarType(ts.ScalarKind.FLOAT64)),
             gtir.Sym(id="size", type=SIZE_TYPE),
         ],
@@ -305,7 +336,7 @@ def test_gtir_cond():
                 expr=im.op_as_fieldop("plus", domain)(
                     "x",
                     im.call("cond")(
-                        gtir.SymRef(id="pred"),
+                        im.greater(gtir.SymRef(id="s1"), gtir.SymRef(id="s2")),
                         im.op_as_fieldop("plus", domain)("y", "scalar"),
                         im.op_as_fieldop("plus", domain)("w", "scalar"),
                     ),
@@ -322,10 +353,10 @@ def test_gtir_cond():
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
 
-    for s in [False, True]:
+    for s1, s2 in [(1, 2), (2, 1)]:
         d = np.empty_like(a)
-        sdfg(pred=np.bool_(s), scalar=1.0, x=a, y=b, w=c, z=d, **FSYMBOLS)
-        assert np.allclose(d, (a + b + 1) if s else (a + c + 1))
+        sdfg(a, b, c, d, s1, s2, scalar=1.0, **FSYMBOLS)
+        assert np.allclose(d, (a + b + 1) if s1 > s2 else (a + c + 1))
 
 
 def test_gtir_cond_nested():
@@ -367,7 +398,7 @@ def test_gtir_cond_nested():
     for s1 in [False, True]:
         for s2 in [False, True]:
             b = np.empty_like(a)
-            sdfg(pred_1=np.bool_(s1), pred_2=np.bool_(s2), x=a, z=b, **FSYMBOLS)
+            sdfg(a, b, pred_1=np.bool_(s1), pred_2=np.bool_(s2), **FSYMBOLS)
             assert np.allclose(b, (a + 1.0) if s1 else (a + 2.0) if s2 else (a + 3.0))
 
 
@@ -468,7 +499,7 @@ def test_gtir_cartesian_shift_left():
 
         FSYMBOLS_tmp = FSYMBOLS.copy()
         FSYMBOLS_tmp["__x_offset_stride_0"] = 1
-        sdfg(x=a, x_offset=a_offset, y=b, **FSYMBOLS_tmp)
+        sdfg(a, a_offset, b, **FSYMBOLS_tmp)
         assert np.allclose(a[OFFSET:] + DELTA, b[:-OFFSET])
 
 
@@ -565,7 +596,7 @@ def test_gtir_cartesian_shift_right():
 
         FSYMBOLS_tmp = FSYMBOLS.copy()
         FSYMBOLS_tmp["__x_offset_stride_0"] = 1
-        sdfg(x=a, x_offset=a_offset, y=b, **FSYMBOLS_tmp)
+        sdfg(a, a_offset, b, **FSYMBOLS_tmp)
         assert np.allclose(a[:-OFFSET] + DELTA, b[OFFSET:])
 
 
@@ -705,8 +736,8 @@ def test_gtir_connectivity_shift():
         ce = np.empty([SIMPLE_MESH.num_cells, SIMPLE_MESH.num_edges])
 
         sdfg(
-            ce_field=ce,
-            ev_field=ev,
+            ce,
+            ev,
             c2e_offset=np.full(SIMPLE_MESH.num_cells, C2E_neighbor_idx, dtype=np.int32),
             e2v_offset=np.full(SIMPLE_MESH.num_edges, E2V_neighbor_idx, dtype=np.int32),
             connectivity_C2E=connectivity_C2E.table,
@@ -777,8 +808,8 @@ def test_gtir_connectivity_shift_chain():
     e_out = np.empty_like(e)
 
     sdfg(
-        edges=e,
-        edges_out=e_out,
+        e,
+        e_out,
         connectivity_E2V=connectivity_E2V.table,
         connectivity_V2E=connectivity_V2E.table,
         **FSYMBOLS,
@@ -837,8 +868,8 @@ def test_gtir_neighbors_as_input():
     ]
 
     sdfg(
-        v2e_field=v2e_field,
-        vertex=v,
+        v2e_field,
+        v,
         **FSYMBOLS,
         **make_mesh_symbols(SIMPLE_MESH),
         __v2e_field_size_0=SIMPLE_MESH.num_vertices,
@@ -895,8 +926,8 @@ def test_gtir_neighbors_as_output():
     v2e_field = np.empty([SIMPLE_MESH.num_vertices, connectivity_V2E.max_neighbors], dtype=e.dtype)
 
     sdfg(
-        edges=e,
-        v2e_field=v2e_field,
+        e,
+        v2e_field,
         connectivity_V2E=connectivity_V2E.table,
         **FSYMBOLS,
         **make_mesh_symbols(SIMPLE_MESH),
@@ -974,8 +1005,8 @@ def test_gtir_reduce():
         v = np.empty(SIMPLE_MESH.num_vertices, dtype=e.dtype)
 
         sdfg(
-            edges=e,
-            vertices=v,
+            e,
+            v,
             connectivity_V2E=connectivity_V2E.table,
             **FSYMBOLS,
             **make_mesh_symbols(SIMPLE_MESH),
@@ -1051,8 +1082,8 @@ def test_gtir_reduce_with_skip_values():
         v = np.empty(SKIP_VALUE_MESH.num_vertices, dtype=e.dtype)
 
         sdfg(
-            edges=e,
-            vertices=v,
+            e,
+            v,
             connectivity_V2E=connectivity_V2E.table,
             **FSYMBOLS,
             **make_mesh_symbols(SKIP_VALUE_MESH),
@@ -1131,8 +1162,8 @@ def test_gtir_reduce_dot_product():
     ]
 
     sdfg(
-        edges=e,
-        vertices=v,
+        e,
+        v,
         connectivity_V2E=connectivity_V2E.table,
         **FSYMBOLS,
         **make_mesh_symbols(SIMPLE_MESH),
@@ -1215,9 +1246,9 @@ def test_gtir_reduce_with_cond_neighbors():
             )
         ]
         sdfg(
-            pred=np.bool_(use_full),
-            edges=e,
-            vertices=v,
+            np.bool_(use_full),
+            e,
+            v,
             connectivity_V2E=connectivity_V2E_skip_values.table,
             connectivity_V2E_FULL=connectivity_V2E_simple.table,
             **FSYMBOLS,
@@ -1274,7 +1305,7 @@ def test_gtir_let_lambda():
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, {})
 
-    sdfg(x=a, y=b, **FSYMBOLS)
+    sdfg(a, b, **FSYMBOLS)
     assert np.allclose(b, a * 8)
 
 
@@ -1314,7 +1345,7 @@ def test_gtir_let_lambda_with_cond():
     a = np.random.rand(N)
     for s in [False, True]:
         b = np.empty_like(a)
-        sdfg(pred=np.bool_(s), x=a, y=b, **FSYMBOLS)
+        sdfg(a, b, pred=np.bool_(s), **FSYMBOLS)
         assert np.allclose(b, a if s else a * 2)
 
 
@@ -1349,5 +1380,5 @@ def test_gtir_if_values():
 
     sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
 
-    sdfg(x=a, y=b, z=c, **FSYMBOLS)
+    sdfg(a, b, c, **FSYMBOLS)
     assert np.allclose(c, np.where(a < b, a, b))