edit

Author: edopao

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

@@ -315,25 +315,18 @@ def translate_broadcast(
     # 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(node.args[0], 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"),