edit

Author: edopao

src/gt4py/next/program_processors/runners/dace/gtir_to_sdfg_primitives.py

@@ -16,7 +16,11 @@
 
 from gt4py.next import common as gtx_common, utils as gtx_utils
 from gt4py.next.iterator import ir as gtir
-from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, domain_utils
+from gt4py.next.iterator.ir_utils import (
+    common_pattern_matcher as cpm,
+    domain_utils,
+    ir_makers as im,
+)
 from gt4py.next.program_processors.runners.dace import (
     gtir_dataflow,
     gtir_domain,
@@ -247,10 +251,16 @@ def translate_as_fieldop(
         raise NotImplementedError("Unexpected 'as_filedop' with tuple output in SDFG lowering.")
 
     if cpm.is_ref_to(fieldop_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(node, ctx, sdfg_builder)
+        if isinstance(node.args[0].type, ts.ScalarType):
+            # Special usage of 'deref' as argument to fieldop expression, to broadcast
+            # a scalar value on the field domain.
+            return translate_broadcast(node, ctx, sdfg_builder)
+        else:
+            # Special usage of 'deref' with field argument, to return a subset of
+            # the full field domain.
+            # TODO(edopao): Lower this case to a memlet edge, planned for next PR.
+            stencil_expr = im.lambda_("a")(im.deref("a"))
+            stencil_expr.expr.type = node.type.dtype
     elif isinstance(fieldop_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.
@@ -285,6 +295,7 @@ def translate_broadcast(
     ctx: gtir_to_sdfg.SubgraphContext,
     sdfg_builder: gtir_to_sdfg.SDFGBuilder,
 ) -> gtir_to_sdfg_types.FieldopData:
+    """Translates a broadcast expression which writes a scalar value on the field domain."""
     assert isinstance(node, gtir.FunCall)
     assert cpm.is_call_to(node.fun, "as_fieldop")
 
@@ -294,6 +305,10 @@ def translate_broadcast(
     assert isinstance(node.type.dtype, ts.ScalarType)
     field_dtype = gtx_dace_utils.as_dace_type(node.type.dtype)
 
+    assert len(node.args) == 1
+    assert isinstance(node.args[0].type, ts.ScalarType)
+    scalar_arg = node.args[0]
+
     fun_node = node.fun
     assert len(fun_node.args) == 2
     fieldop_expr, fieldop_domain_expr = fun_node.args
@@ -308,32 +323,23 @@ def translate_broadcast(
     # The memory layout of the output field follows the field operator compute domain.
     field_dims, field_origin, field_shape = gtir_domain.get_field_layout(field_domain)
     assert field_dims == node.type.dims
-
     field_name, field_desc = sdfg_builder.add_temp_array(ctx.sdfg, field_shape, field_dtype)
     field_node = ctx.state.add_access(field_name)
 
     # Retrieve the scalar argument, which could be either a literal value or the
     # result of a scalar expression.
-    assert len(node.args) == 1
+    arg = _parse_fieldop_arg(scalar_arg, ctx, sdfg_builder, field_domain)
+    assert isinstance(arg, gtir_dataflow.MemletExpr)
+    assert arg.subset.num_elements() == 1
 
     # Use a 'Fill' library node to write the scalar value to the result field.
-    if isinstance(node.args[0], gtir.Literal):
-        assert node.args[0].type == node.type.dtype
-        value = field_dtype(node.args[0].value)
-        fill_node = sdfg_library_nodes.Fill("fill", value)
-        ctx.state.add_node(fill_node)
-    else:
-        arg = _parse_fieldop_arg(node.args[0], ctx, sdfg_builder, field_domain)
-        assert isinstance(arg, gtir_dataflow.MemletExpr)
-        assert arg.subset.num_elements() == 1
-
-        fill_node = sdfg_library_nodes.Fill("fill")
-        ctx.state.add_node(fill_node)
-        ctx.state.add_nedge(
-            arg.dc_node,
-            fill_node,
-            dace.Memlet(data=arg.dc_node.data, subset=arg.subset),
-        )
+    fill_node = sdfg_library_nodes.Fill("fill")
+    ctx.state.add_node(fill_node)
+    ctx.state.add_nedge(
+        arg.dc_node,
+        fill_node,
+        dace.Memlet(data=arg.dc_node.data, subset=arg.subset),
+    )
 
     ctx.state.add_nedge(
         fill_node,

src/gt4py/next/program_processors/runners/dace/sdfg_library_nodes.py

@@ -44,20 +44,16 @@ def expansion(node: Fill, parent_state: dace.SDFGState, parent_sdfg: dace.SDFG)
         out_mem = dace.Memlet(expr=f"{out}[{','.join(map_params)}]")
         outputs = {"_out": out_mem}
 
-        if node._value is None:
-            assert len(parent_state.in_edges(node)) == 1
-            inedge = parent_state.in_edges(node)[0]
-            inp_desc = parent_sdfg.arrays[inedge.data.data]
-            inner_inp_desc = inp_desc.clone()
-            inner_inp_desc.transient = False
-            inp = sdfg.add_datadesc(_INPUT_NAME, inner_inp_desc)
-            inedge._dst_conn = _INPUT_NAME
-            node.add_in_connector(_INPUT_NAME)
-            inputs = {"_in": dace.Memlet(data=inp, subset="0")}
-            code = "_out = _in"
-        else:
-            inputs = {}
-            code = f"_out = {node._value}"
+        assert len(parent_state.in_edges(node)) == 1
+        inedge = parent_state.in_edges(node)[0]
+        inp_desc = parent_sdfg.arrays[inedge.data.data]
+        inner_inp_desc = inp_desc.clone()
+        inner_inp_desc.transient = False
+        inp = sdfg.add_datadesc(_INPUT_NAME, inner_inp_desc)
+        inedge._dst_conn = _INPUT_NAME
+        node.add_in_connector(_INPUT_NAME)
+        inputs = {"_in": dace.Memlet(data=inp, subset="0")}
+        code = "_out = _in"
 
         state.add_mapped_tasklet(
             f"{node.label}_tasklet", map_rng, inputs, code, outputs, external_edges=True
@@ -72,8 +68,6 @@ class Fill(dace_nodes.LibraryNode):
 
     implementations: Final[dict[str, dace_transform.ExpandTransformation]] = {"pure": ExpandPure}
     default_implementation: Final[str] = "pure"
-    _value: dace.typeclass | None
 
-    def __init__(self, name: str, value: dace.typeclass | None = None):
+    def __init__(self, name: str):
         super().__init__(name)
-        self._value = value

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py

@@ -2110,16 +2110,8 @@ def test_gtir_concat_where():
         gtx_common.GridType.CARTESIAN, {IDim: (SUBSET_SIZE, gtir.InfinityLiteral.POSITIVE)}
     )
 
-    concat_expr_lhs = im.concat_where(
-        domain_cond_lhs,
-        im.as_fieldop("deref")("x"),
-        im.as_fieldop("deref")("y"),
-    )
-    concat_expr_rhs = im.concat_where(
-        domain_cond_rhs,
-        im.as_fieldop("deref")("y"),
-        im.as_fieldop("deref")("x"),
-    )
+    concat_expr_lhs = im.concat_where(domain_cond_lhs, "x", "y")
+    concat_expr_rhs = im.concat_where(domain_cond_rhs, "y", "x")
 
     a = np.random.rand(N)
     b = np.random.rand(N)
@@ -2177,12 +2169,8 @@ def test_gtir_concat_where_two_dimensions():
             gtir.SetAt(
                 expr=im.concat_where(
                     domain_cond1,  # 0, 30; 10,20
-                    im.concat_where(
-                        domain_cond2,
-                        im.as_fieldop("deref")("x"),
-                        im.as_fieldop("deref")("y"),
-                    ),
-                    im.as_fieldop("deref")("w"),
+                    im.concat_where(domain_cond2, "x", "y"),
+                    "w",
                 ),
                 domain=domain,
                 target=gtir.SymRef(id="z"),