Repository: GridTools/gt4py
Pull request #1918 (2025-03-17 18:20:47)

feat[next]: GTIR broadcast builtin
Previously, a broadcast generated an `as_fieldop`. In [PR#1853](https://github.com/GridTools/gt4py/pull/1853), the domain restriction returns a smaller domain for the `as_fieldop`.
In order to avoid this, we introduce the broadcast builtin here and transform it into an `im.as_fieldop("deref", ...)` call after the domain inference.


Patch:
diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py
index f884ec55..d76bc250 100644
--- a/src/gt4py/next/ffront/foast_to_gtir.py
+++ b/src/gt4py/next/ffront/foast_to_gtir.py
@@ -224,9 +224,18 @@ class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator):
     def visit_Symbol(self, node: foast.Symbol, **kwargs: Any) -> itir.Sym:
         return im.sym(node.id)
 
-    def visit_Name(self, node: foast.Name, **kwargs: Any) -> itir.SymRef:
+    def visit_Name(self, node: foast.Name, **kwargs: Any) -> itir.SymRef | itir.AxisLiteral:
+        if isinstance(node.type, ts.DimensionType):
+            return itir.AxisLiteral(value=node.type.dim.value, kind=node.type.dim.kind)
         return im.ref(node.id)
 
+    def visit_Attribute(self, node: foast.Attribute, **kwargs: Any) -> itir.AxisLiteral:
+        if isinstance(node.type, ts.DimensionType):
+            return itir.AxisLiteral(value=node.type.dim.value, kind=node.type.dim.kind)
+        raise AssertionError(
+            "Unexpected attribute access. At this point all accesses should have been removed by `ClosureVarFolding`."
+        )
+
     def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> itir.Expr:
         return im.tuple_get(node.index, self.visit(node.value, **kwargs))
 
@@ -397,8 +406,7 @@ class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator):
     _visit_concat_where = _visit_where  # TODO(havogt): upgrade concat_where
 
     def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
-        expr = self.visit(node.args[0], **kwargs)
-        return im.as_fieldop(im.ref("deref"))(expr)
+        return im.call("broadcast")(*self.visit(node.args, **kwargs))
 
     def _visit_math_built_in(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         return self._lower_and_map(self.visit(node.func, **kwargs), *node.args)
diff --git a/src/gt4py/next/iterator/builtins.py b/src/gt4py/next/iterator/builtins.py
index 8e5f7add..9b065de6 100644
--- a/src/gt4py/next/iterator/builtins.py
+++ b/src/gt4py/next/iterator/builtins.py
@@ -27,6 +27,11 @@ def index(*args):
     raise BackendNotSelectedError()
 
 
+@builtin_dispatch
+def broadcast(*args):
+    raise BackendNotSelectedError()
+
+
 @builtin_dispatch
 def deref(*args):
     raise BackendNotSelectedError()
@@ -470,6 +475,7 @@ ARITHMETIC_BUILTINS = {
 
 BUILTINS = {
     "as_fieldop",  # `as_fieldop(stencil, domain)` creates field_operator from stencil (domain is optional, but for now required for embedded execution)
+    "broadcast",
     "can_deref",
     "cartesian_domain",
     "cast_",
diff --git a/src/gt4py/next/iterator/ir_utils/misc.py b/src/gt4py/next/iterator/ir_utils/misc.py
index 03652cdf..d604d976 100644
--- a/src/gt4py/next/iterator/ir_utils/misc.py
+++ b/src/gt4py/next/iterator/ir_utils/misc.py
@@ -71,3 +71,10 @@ def is_equal(a: itir.Expr, b: itir.Expr):
     return a == b or (
         CannonicalizeBoundSymbolNames.apply(a) == CannonicalizeBoundSymbolNames.apply(b)
     )
+
+
+def with_altered_arg(node: itir.FunCall, arg_idx: int, new_arg: itir.Expr | str) -> itir.FunCall:
+    """Given a itir.FunCall return a new call with one of its argument replaced."""
+    return im.call(node.fun)(
+        *(arg if i != arg_idx else im.ensure_expr(new_arg) for i, arg in enumerate(node.args))
+    )
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
index 462f87b6..7a10587d 100644
--- a/src/gt4py/next/iterator/transforms/collapse_tuple.py
+++ b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -29,14 +29,6 @@ from gt4py.next.iterator.type_system import inference as itir_type_inference
 from gt4py.next.type_system import type_info, type_specifications as ts
 
 
-def _with_altered_arg(node: ir.FunCall, arg_idx: int, new_arg: ir.Expr | str):
-    """Given a itir.FunCall return a new call with one of its argument replaced."""
-    return ir.FunCall(
-        fun=node.fun,
-        args=[arg if i != arg_idx else im.ensure_expr(new_arg) for i, arg in enumerate(node.args)],
-    )
-
-
 def _is_trivial_make_tuple_call(node: ir.Expr):
     """Return if node is a `make_tuple` call with all elements `SymRef`s, `Literal`s or tuples thereof."""
     if not cpm.is_call_to(node, "make_tuple"):
@@ -334,10 +326,10 @@ class CollapseTuple(
                 if cpm.is_call_to(arg, "if_"):
                     cond, true_branch, false_branch = arg.args
                     new_true_branch = self.fp_transform(
-                        _with_altered_arg(node, i, true_branch), **kwargs
+                        ir_misc.with_altered_arg(node, i, true_branch), **kwargs
                     )
                     new_false_branch = self.fp_transform(
-                        _with_altered_arg(node, i, false_branch), **kwargs
+                        ir_misc.with_altered_arg(node, i, false_branch), **kwargs
                     )
                     return im.if_(cond, new_true_branch, new_false_branch)
         return None
@@ -397,7 +389,7 @@ class CollapseTuple(
                     for type_ in tuple_type.types
                 ]
                 f_args = [im.ref(param.id, param.type) for param in f_params]
-                f_body = _with_altered_arg(node, i, im.make_tuple(*f_args))
+                f_body = ir_misc.with_altered_arg(node, i, im.make_tuple(*f_args))
                 # simplify, e.g., inline trivial make_tuple args
                 new_f_body = self.fp_transform(f_body, **kwargs)
                 # if the continuation did not simplify there is nothing to gain. Skip
diff --git a/src/gt4py/next/iterator/transforms/infer_domain.py b/src/gt4py/next/iterator/transforms/infer_domain.py
index f3c31852..239824ba 100644
--- a/src/gt4py/next/iterator/transforms/infer_domain.py
+++ b/src/gt4py/next/iterator/transforms/infer_domain.py
@@ -20,6 +20,7 @@ from gt4py.next.iterator.ir_utils import (
     common_pattern_matcher as cpm,
     domain_utils,
     ir_makers as im,
+    misc as ir_misc,
 )
 from gt4py.next.iterator.transforms import trace_shifts
 from gt4py.next.utils import flatten_nested_tuple, tree_map
@@ -363,6 +364,19 @@ def _infer_if(
     return result_expr, actual_domains
 
 
+def _infer_broadcast(
+    expr: itir.Expr,
+    domain: DomainAccess,
+    **kwargs: Unpack[InferenceOptions],
+) -> tuple[itir.Expr, AccessedDomains]:
+    assert cpm.is_call_to(expr, "broadcast")
+    # We just propagate the domain to the first argument. Restriction of the domain is based
+    # on the type and occurs in a general setting (not yet merged #1853).
+    infered_expr, actual_domains = infer_expr(expr.args[0], domain, **kwargs)
+
+    return ir_misc.with_altered_arg(expr, 0, infered_expr), actual_domains
+
+
 def _infer_expr(
     expr: itir.Expr,
     domain: DomainAccess,
@@ -382,6 +396,8 @@ def _infer_expr(
         return _infer_tuple_get(expr, domain, **kwargs)
     elif cpm.is_call_to(expr, "if_"):
         return _infer_if(expr, domain, **kwargs)
+    elif cpm.is_call_to(expr, "broadcast"):
+        return _infer_broadcast(expr, domain, **kwargs)
     elif (
         cpm.is_call_to(expr, builtins.ARITHMETIC_BUILTINS)
         or cpm.is_call_to(expr, builtins.TYPE_BUILTINS)
diff --git a/src/gt4py/next/iterator/transforms/pass_manager.py b/src/gt4py/next/iterator/transforms/pass_manager.py
index 4023950d..19b13876 100644
--- a/src/gt4py/next/iterator/transforms/pass_manager.py
+++ b/src/gt4py/next/iterator/transforms/pass_manager.py
@@ -18,6 +18,7 @@ from gt4py.next.iterator.transforms import (
     inline_dynamic_shifts,
     inline_fundefs,
     inline_lifts,
+    prune_broadcast,
 )
 from gt4py.next.iterator.transforms.collapse_list_get import CollapseListGet
 from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
@@ -92,6 +93,7 @@ def apply_common_transforms(
         offset_provider=offset_provider,
         symbolic_domain_sizes=symbolic_domain_sizes,
     )
+    ir = prune_broadcast.PruneBroadcast.apply(ir)
 
     for _ in range(10):
         inlined = ir
@@ -187,4 +189,5 @@ def apply_fieldview_transforms(
         ir
     )  # domain inference does not support dynamic offsets yet
     ir = infer_domain.infer_program(ir, offset_provider=offset_provider)
+    ir = prune_broadcast.PruneBroadcast.apply(ir)
     return ir
diff --git a/src/gt4py/next/iterator/transforms/prune_broadcast.py b/src/gt4py/next/iterator/transforms/prune_broadcast.py
new file mode 100644
index 00000000..432e9d2d
--- /dev/null
+++ b/src/gt4py/next/iterator/transforms/prune_broadcast.py
@@ -0,0 +1,35 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+import dataclasses
+
+from gt4py.eve import NodeTranslator, PreserveLocationVisitor
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import (
+    common_pattern_matcher as cpm,
+    domain_utils,
+    ir_makers as im,
+)
+
+
+@dataclasses.dataclass
+class PruneBroadcast(PreserveLocationVisitor, NodeTranslator):
+    PRESERVED_ANNEX_ATTRS = ("domain",)
+
+    @classmethod
+    def apply(cls, node: itir.Node):
+        return cls().visit(node)
+
+    def visit_FunCall(self, node: itir.FunCall) -> itir.FunCall:
+        node = self.generic_visit(node)
+
+        if cpm.is_call_to(node, "broadcast"):
+            expr = self.visit(node.args[0])
+            node = im.as_fieldop("deref", domain_utils.SymbolicDomain.as_expr(node.annex.domain))(
+                expr
+            )
+        return node
diff --git a/src/gt4py/next/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py
index 131b773d..df36189b 100644
--- a/src/gt4py/next/iterator/type_system/type_synthesizer.py
+++ b/src/gt4py/next/iterator/type_system/type_synthesizer.py
@@ -205,6 +205,25 @@ def index(arg: ts.DimensionType) -> ts.FieldType:
     )
 
 
+@_register_builtin_type_synthesizer
+def broadcast(
+    arg: ts.FieldType | ts.ScalarType | ts.DeferredType, dims: tuple[ts.DimensionType]
+) -> ts.FieldType | ts.DeferredType:
+    if isinstance(arg, ts.DeferredType):
+        return arg
+
+    dims_ = [dim.dim for dim in dims]
+
+    if isinstance(arg, ts.FieldType):
+        dtype = arg.dtype
+    elif isinstance(arg, ts.ScalarType):
+        dtype = arg
+    else:
+        raise AssertionError("`args` need to be either `ts.FieldType` or `ts.ScalarType`.")
+
+    return ts.FieldType(dims=dims_, dtype=dtype)
+
+
 @_register_builtin_type_synthesizer
 def neighbors(offset_literal: it_ts.OffsetLiteralType, it: it_ts.IteratorType) -> ts.ListType:
     assert (


Test patch:
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py
index 776cd4e1..53661427 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py
@@ -38,7 +38,7 @@ from gt4py.next.ffront.func_to_foast import FieldOperatorParser
 from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms import inline_lambdas
 from gt4py.next.type_system import type_specifications as ts, type_translation
-
+from gt4py.next.iterator import ir as itir
 
 Edge = gtx.Dimension("Edge")
 Vertex = gtx.Dimension("Vertex")
@@ -907,7 +907,10 @@ def test_broadcast():
     lowered = FieldOperatorLowering.apply(parsed)
 
     assert lowered.id == "foo"
-    assert lowered.expr == im.as_fieldop("deref")(im.ref("inp"))
+    assert lowered.expr == im.call("broadcast")(
+        im.ref("inp"),
+        im.make_tuple(*(itir.AxisLiteral(value=dim.value, kind=dim.kind) for dim in (TDim, UDim))),
+    )
 
 
 def test_scalar_broadcast():
@@ -918,4 +921,7 @@ def test_scalar_broadcast():
     lowered = FieldOperatorLowering.apply(parsed)
 
     assert lowered.id == "foo"
-    assert lowered.expr == im.as_fieldop("deref")(1)
+    assert lowered.expr == im.call("broadcast")(
+        1,
+        im.make_tuple(*(itir.AxisLiteral(value=dim.value, kind=dim.kind) for dim in (TDim, UDim))),
+    )
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py
index 86cc8a67..4385cbe7 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py
@@ -1093,3 +1093,12 @@ def test_never_accessed_domain_tuple(offset_provider):
         "in_field2": infer_domain.DomainAccessDescriptor.NEVER,
     }
     run_test_expr(testee, testee, domain, expected_domains, offset_provider)
+
+
+def test_broadcast(offset_provider):
+    testee = im.call("broadcast")("in_field", im.make_tuple(itir.AxisLiteral(value="IDim")))
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 10)})
+    expected_domains = {
+        "in_field": {IDim: (0, 10)},
+    }
+    run_test_expr(testee, testee, domain, expected_domains, offset_provider)