Repository: GridTools/gt4py
Pull request #1693 (2024-10-17 16:14:04)

feat[next][dace]: GTIR-to-SDFG broadcast tuple scalar arg + cleanup
This PR implements the lowering of expressions like:
`as_fieldop(deref, domain)(x[0] + (x[1][0] × 2.0 + x[1][1] × 3.0))`
where the elements of tuple `x` being accessed are all scalar values, so `x[0] + (x[1][0] × 2.0 + x[1][1] × 3.0)` is a scalar expression.

Therefore, this PR allows to lower expressions where the argument to `as_field_op` is a scalar expression rather than a lambda. Note that:
- in case of a lambda, the `as_fieldop` function computes the lambda over the field domain;
- in case of scalar expression, it broadcasts the scalar value over the field domain.

Additionally, some cleanup:
- Removed call to `gt_simplify` so that dace backend returns a plain unoptimized SDFG.
- Removed `patch_gtir`, since the functionality of this IR pass is provided by similar changes delivered in #1677 and #1689. 


Patch:
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 8fb1451e..a8ae1cc0 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
@@ -209,8 +209,8 @@ def translate_as_fieldop(
     The dataflow can be as simple as a single tasklet, or implement a local computation
     as a composition of tasklets and even include a map to range on local dimensions (e.g.
     neighbors and map builtins).
-    The stencil dataflow is instantiated inside a map scope, which apply the stencil over
-    the field domain.
+    The stencil dataflow is instantiated inside a map scope, which applies the stencil
+    over the field domain.
     """
     assert isinstance(node, gtir.FunCall)
     assert cpm.is_call_to(node.fun, "as_fieldop")
@@ -219,11 +219,24 @@ def translate_as_fieldop(
     fun_node = node.fun
     assert len(fun_node.args) == 2
     stencil_expr, domain_expr = fun_node.args
-    assert isinstance(stencil_expr, gtir.Lambda)
-    assert isinstance(domain_expr, gtir.FunCall)
+
+    if isinstance(stencil_expr, gtir.Lambda):
+        # Default case, handled below: the argument expression is a lambda function
+        # representing the stencil operation to be computed over the field domain.
+        pass
+    elif cpm.is_ref_to(stencil_expr, "deref"):
+        # Special usage of 'deref' as argument to fieldop expression, to pass a scalar
+        # value to 'as_fieldop' function. It results in broadcasting the scalar value
+        # over the field domain.
+        return translate_broadcast_scalar(node, sdfg, state, sdfg_builder, reduce_identity)
+    else:
+        raise NotImplementedError(
+            f"Expression type '{type(stencil_expr)}' not supported as argument to 'as_fieldop' node."
+        )
 
     # parse the domain of the field operator
     domain = extract_domain(domain_expr)
+    domain_indices = sbs.Indices([dace_gtir_utils.get_map_variable(dim) for dim, _, _ in domain])
 
     # The reduction identity value is used in place of skip values when building
     # a list of neighbor values in the unstructured domain.
@@ -273,16 +286,15 @@ def translate_as_fieldop(
     input_edges, output = taskgen.visit(stencil_expr, args=stencil_args)
     output_desc = output.result.node.desc(sdfg)
 
-    domain_index = sbs.Indices([dace_gtir_utils.get_map_variable(dim) for dim, _, _ in domain])
     if isinstance(node.type.dtype, itir_ts.ListType):
         assert isinstance(output_desc, dace.data.Array)
         assert set(output_desc.offset) == {0}
         # additional local dimension for neighbors
         # TODO(phimuell): Investigate if we should swap the two.
-        output_subset = sbs.Range.from_indices(domain_index) + sbs.Range.from_array(output_desc)
+        output_subset = sbs.Range.from_indices(domain_indices) + sbs.Range.from_array(output_desc)
     else:
         assert isinstance(output_desc, dace.data.Scalar)
-        output_subset = sbs.Range.from_indices(domain_index)
+        output_subset = sbs.Range.from_indices(domain_indices)
 
     # create map range corresponding to the field operator domain
     me, mx = sdfg_builder.add_map(
@@ -307,6 +319,62 @@ def translate_as_fieldop(
     return result_field
 
 
+def translate_broadcast_scalar(
+    node: gtir.Node,
+    sdfg: dace.SDFG,
+    state: dace.SDFGState,
+    sdfg_builder: gtir_sdfg.SDFGBuilder,
+    reduce_identity: Optional[gtir_dataflow.SymbolExpr],
+) -> FieldopResult:
+    """
+    Generates the dataflow subgraph for the 'as_fieldop' builtin function for the
+    special case where the argument to 'as_fieldop' is a 'deref' scalar expression,
+    rather than a lambda function. This case corresponds to broadcasting the scalar
+    value over the field domain. Therefore, it is lowered to a mapped tasklet that
+    just writes the scalar value out to all elements of the result field.
+    """
+    assert isinstance(node, gtir.FunCall)
+    assert cpm.is_call_to(node.fun, "as_fieldop")
+    assert isinstance(node.type, ts.FieldType)
+
+    fun_node = node.fun
+    assert len(fun_node.args) == 2
+    stencil_expr, domain_expr = fun_node.args
+    assert cpm.is_ref_to(stencil_expr, "deref")
+
+    domain = extract_domain(domain_expr)
+    domain_indices = sbs.Indices([dace_gtir_utils.get_map_variable(dim) for dim, _, _ in domain])
+
+    assert len(node.args) == 1
+    assert isinstance(node.args[0].type, ts.ScalarType)
+    scalar_expr = _parse_fieldop_arg(
+        node.args[0], sdfg, state, sdfg_builder, domain, reduce_identity=None
+    )
+    assert isinstance(scalar_expr, gtir_dataflow.MemletExpr)
+    assert scalar_expr.subset == sbs.Indices.from_string("0")
+    result = gtir_dataflow.DataflowOutputEdge(
+        state, gtir_dataflow.DataExpr(scalar_expr.node, node.args[0].type)
+    )
+    result_field = _create_temporary_field(sdfg, state, domain, node.type, dataflow_output=result)
+
+    sdfg_builder.add_mapped_tasklet(
+        "broadcast",
+        state,
+        map_ranges={
+            dace_gtir_utils.get_map_variable(dim): f"{lower_bound}:{upper_bound}"
+            for dim, lower_bound, upper_bound in domain
+        },
+        inputs={"__inp": dace.Memlet(data=scalar_expr.node.data, subset="0")},
+        code="__val = __inp",
+        outputs={"__val": dace.Memlet(data=result_field.data_node.data, subset=domain_indices)},
+        input_nodes={scalar_expr.node.data: scalar_expr.node},
+        output_nodes={result_field.data_node.data: result_field.data_node},
+        external_edges=True,
+    )
+
+    return result_field
+
+
 def translate_if(
     node: gtir.Node,
     sdfg: dace.SDFG,
@@ -638,6 +706,7 @@ if TYPE_CHECKING:
     # Use type-checking to assert that all translator functions implement the `PrimitiveTranslator` protocol
     __primitive_translators: list[PrimitiveTranslator] = [
         translate_as_fieldop,
+        translate_broadcast_scalar,
         translate_if,
         translate_literal,
         translate_make_tuple,
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
index d79d8873..3697609d 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
@@ -32,7 +32,6 @@ from gt4py.next.program_processors.runners.dace_common import utility as dace_ut
 from gt4py.next.program_processors.runners.dace_fieldview import (
     gtir_builtin_translators,
     gtir_dataflow,
-    transformations as gtx_transformations,
     utility as dace_gtir_utils,
 )
 from gt4py.next.type_system import type_specifications as ts, type_translation as tt
@@ -665,7 +664,6 @@ def build_sdfg_from_gtir(
 
     The lowering to SDFG requires that the program node is type-annotated, therefore this function
     runs type ineference as first step.
-    As a final step, it runs the `simplify` pass to ensure that the SDFG is in the DaCe canonical form.
 
     Arguments:
         ir: The GTIR program node to be lowered to SDFG
@@ -677,10 +675,8 @@ def build_sdfg_from_gtir(
 
     ir = gtir_type_inference.infer(ir, offset_provider=offset_provider)
     ir = ir_prune_casts.PruneCasts().visit(ir)
-    ir = dace_gtir_utils.patch_gtir(ir)
     sdfg_genenerator = GTIRToSDFG(offset_provider)
     sdfg = sdfg_genenerator.visit(ir)
     assert isinstance(sdfg, dace.SDFG)
 
-    gtx_transformations.gt_simplify(sdfg)
     return sdfg
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 855dc9c9..b5c447a1 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
@@ -11,10 +11,7 @@ from __future__ import annotations
 import itertools
 from typing import Any
 
-from gt4py import eve
 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.type_system import type_specifications as ts
 
 
@@ -65,41 +62,3 @@ def get_tuple_type(data: tuple[Any, ...]) -> ts.TupleType:
     return ts.TupleType(
         types=[get_tuple_type(d) if isinstance(d, tuple) else d.data_type for d in data]
     )
-
-
-def patch_gtir(ir: gtir.Program) -> gtir.Program:
-    """
-    Make the IR compliant with the requirements of lowering to SDFG.
-
-    Applies canonicalization of as_fieldop expressions as well as some temporary workarounds.
-    This allows to lower the IR to SDFG for some special cases.
-    """
-
-    class PatchGTIR(eve.PreserveLocationVisitor, eve.NodeTranslator):
-        def visit_FunCall(self, node: gtir.FunCall) -> gtir.Node:
-            if cpm.is_applied_as_fieldop(node):
-                assert isinstance(node.fun, gtir.FunCall)
-                assert isinstance(node.type, ts.FieldType)
-
-                # Handle the case of fieldop without domain. This case should never happen, but domain
-                # inference currently produces this kind of nodes for unreferenced tuple fields.
-                # TODO(tehrengruber): remove this workaround once domain ineference supports this case
-                if len(node.fun.args) == 1:
-                    return gtir.Literal(value="0", type=node.type.dtype)
-
-                assert len(node.fun.args) == 2
-                stencil = node.fun.args[0]
-
-                # Canonicalize as_fieldop: always expect a lambda expression.
-                # Here we replace the call to deref with a lambda expression and empty arguments list.
-                if cpm.is_ref_to(stencil, "deref"):
-                    node.fun.args[0] = gtir.Lambda(
-                        expr=gtir.FunCall(fun=stencil, args=node.args), params=[]
-                    )
-                    node.args = []
-
-            node.args = self.visit(node.args)
-            node.fun = self.visit(node.fun)
-            return node
-
-    return PatchGTIR().visit(ir)


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 98e15dac..728b4b02 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
@@ -40,10 +40,11 @@ dace_backend = pytest.importorskip("gt4py.next.program_processors.runners.dace_f
 
 
 N = 10
-IFTYPE = ts.FieldType(dims=[IDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-CFTYPE = ts.FieldType(dims=[Cell], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-EFTYPE = ts.FieldType(dims=[Edge], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-VFTYPE = ts.FieldType(dims=[Vertex], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
+FLOAT_TYPE = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+IFTYPE = ts.FieldType(dims=[IDim], dtype=FLOAT_TYPE)
+CFTYPE = ts.FieldType(dims=[Cell], dtype=FLOAT_TYPE)
+EFTYPE = ts.FieldType(dims=[Edge], dtype=FLOAT_TYPE)
+VFTYPE = ts.FieldType(dims=[Vertex], dtype=FLOAT_TYPE)
 V2E_FTYPE = ts.FieldType(dims=[Vertex, V2EDim], dtype=EFTYPE.dtype)
 CARTESIAN_OFFSETS = {
     "IDim": IDim,
@@ -315,6 +316,62 @@ def test_gtir_tuple_expr():
     assert np.allclose(c, a * 2 + b)
 
 
+def test_gtir_tuple_broadcast_scalar():
+    domain = im.domain(gtx_common.GridType.CARTESIAN, ranges={IDim: (0, "size")})
+    testee = gtir.Program(
+        id="gtir_tuple_broadcast_scalar",
+        function_definitions=[],
+        params=[
+            gtir.Sym(
+                id="x",
+                type=ts.TupleType(types=[FLOAT_TYPE, ts.TupleType(types=[FLOAT_TYPE, FLOAT_TYPE])]),
+            ),
+            gtir.Sym(id="y", type=IFTYPE),
+            gtir.Sym(id="size", type=SIZE_TYPE),
+        ],
+        declarations=[],
+        body=[
+            gtir.SetAt(
+                expr=im.as_fieldop("deref", domain)(
+                    im.plus(
+                        im.tuple_get(0, "x"),
+                        im.plus(
+                            im.multiplies_(
+                                im.tuple_get(
+                                    0,
+                                    im.tuple_get(1, "x"),
+                                ),
+                                2.0,
+                            ),
+                            im.multiplies_(
+                                im.tuple_get(
+                                    1,
+                                    im.tuple_get(1, "x"),
+                                ),
+                                3.0,
+                            ),
+                        ),
+                    )
+                ),
+                domain=domain,
+                target=gtir.SymRef(id="y"),
+            )
+        ],
+    )
+
+    a = np.random.rand()
+    b = np.random.rand()
+    c = np.random.rand()
+    d = np.empty(N, dtype=type(a))
+
+    sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
+
+    x_fields = (a, b, c)
+
+    sdfg(*x_fields, d, **FSYMBOLS)
+    assert np.allclose(d, a + 2 * b + 3 * c)
+
+
 def test_gtir_tuple_return():
     domain = im.domain(gtx_common.GridType.CARTESIAN, ranges={IDim: (0, "size")})
     testee = gtir.Program(
@@ -954,8 +1011,8 @@ def test_gtir_connectivity_shift():
         im.op_as_fieldop("plus", edge_domain)("e2v_offset", 0),
     )
 
-    CE_FTYPE = ts.FieldType(dims=[Cell, Edge], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-    EV_FTYPE = ts.FieldType(dims=[Edge, Vertex], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
+    CE_FTYPE = ts.FieldType(dims=[Cell, Edge], dtype=FLOAT_TYPE)
+    EV_FTYPE = ts.FieldType(dims=[Edge, Vertex], dtype=FLOAT_TYPE)
     CELL_OFFSET_FTYPE = ts.FieldType(dims=[Cell], dtype=SIZE_TYPE)
     EDGE_OFFSET_FTYPE = ts.FieldType(dims=[Edge], dtype=SIZE_TYPE)