feat[next]: Integration of concat_where (#1713)

In previous commits the following pieces where already added:
- #1998 
- #2065 

In this PR, `concat_where` tests are added and transformations are
integrated. Currently, the feature is disabled for embedded (#2127) and
DaCe (PR will be opened once this is merged).

Details:
- constant folding for `InfinityLiteral`s is added
- various tuning to passes for `index`, `concat_where` builtins and
`InfinityLiteral`
- cases: field constructors now accept `Domain`s (instead of
`dict[Dimension, int])

---------

Co-authored-by: Hannes Vogt <[email protected]>
Co-authored-by: Till Ehrengruber <[email protected]>
Co-authored-by: Edoardo Paone <[email protected]>

Author: 4 people

src/gt4py/next/ffront/foast_to_gtir.py

@@ -406,7 +406,9 @@ def create_if(
 
         return im.let(cond_symref_name, cond_)(result)
 
-    _visit_concat_where = _visit_where  # TODO(havogt): upgrade concat_where
+    def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
+        domain, true_branch, false_branch = self.visit(node.args, **kwargs)
+        return im.concat_where(domain, true_branch, false_branch)
 
     def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         return im.call("broadcast")(*self.visit(node.args, **kwargs))
@@ -488,7 +490,7 @@ def _map(
     Mapping includes making the operation an `as_fieldop` (first kind of mapping), but also `itir.map_`ing lists.
     """
     if all(
-        isinstance(t, ts.ScalarType)
+        isinstance(t, (ts.ScalarType, ts.DimensionType, ts.DomainType))
         for arg_type in original_arg_types
         for t in type_info.primitive_constituents(arg_type)
     ):

src/gt4py/next/iterator/builtins.py

@@ -407,6 +407,11 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     raise BackendNotSelectedError()
 
 
+@builtin_dispatch
+def concat_where(*args):
+    raise BackendNotSelectedError()
+
+
 UNARY_MATH_NUMBER_BUILTINS = {"abs", "neg"}
 UNARY_LOGICAL_BUILTINS = {"not_"}
 UNARY_MATH_FP_BUILTINS = {
@@ -494,6 +499,7 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     "scan",
     "tuple_get",
     "unstructured_domain",
+    "concat_where",
     *ARITHMETIC_BUILTINS,
     *TYPE_BUILTINS,
 }

src/gt4py/next/iterator/embedded.py

@@ -1802,6 +1802,11 @@ def index(axis: common.Dimension) -> common.Field:
     return IndexField(axis)
 
 
+@builtins.concat_where.register(EMBEDDED)
+def concat_where(*args):
+    raise NotImplementedError("To be implemented in frontend embedded.")
+
+
 def closure(
     domain_: runtime.CartesianDomain | runtime.UnstructuredDomain,
     sten: Callable[..., Any],

src/gt4py/next/iterator/transforms/collapse_tuple.py

@@ -50,7 +50,7 @@ def _with_altered_iterator_position_dims(
     )
 
 
-def _is_trivial_make_tuple_call(node: ir.Expr):
+def _is_trivial_make_tuple_call(node: itir.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"):
         return False
@@ -307,9 +307,10 @@ def transform_propagate_tuple_get(self, node: itir.FunCall, **kwargs) -> Optiona
                         self.fp_transform(im.tuple_get(idx.value, expr.fun.expr), **kwargs)
                     )
                 )(*expr.args)
-            elif cpm.is_call_to(expr, "if_"):
+            elif cpm.is_call_to(expr, ("if_", "concat_where")):
+                fun = expr.fun
                 cond, true_branch, false_branch = expr.args
-                return im.if_(
+                return im.call(fun)(
                     cond,
                     self.fp_transform(im.tuple_get(idx.value, true_branch), **kwargs),
                     self.fp_transform(im.tuple_get(idx.value, false_branch), **kwargs),

src/gt4py/next/iterator/transforms/concat_where/transform_to_as_fieldop.py

@@ -25,8 +25,8 @@ def _in(pos: itir.Expr, domain: itir.Expr) -> itir.Expr:
     """
     Given a position and a domain return an expression that evaluates to `True` if the position is inside the domain.
 
-    `in_({i, j, k}, u⟨ Iₕ: [i0, i1[, Iₕ: [j0, j1[, Iₕ: [k0, k1[ ⟩`
-    -> `i0 <= i < i1 & j0 <= j < j1 & k0 <= k < k1`
+    pos = `{i, j, k}`, domain = `u⟨ Iₕ: [i0, i1[, Iₕ: [j0, j1[, Iₕ: [k0, k1[ ⟩`
+    -> `((i0 <= i) & (i < i1)) & ((j0 <= j) & (j < j1)) & ((k0 <= k)l & (k < k1))`
     """
     ret = [
         im.and_(

src/gt4py/next/iterator/transforms/constant_folding.py

@@ -98,6 +98,8 @@ class Transformation(enum.Flag):
         # `if_(True, true_branch, false_branch)` -> `true_branch`
         FOLD_IF = enum.auto()
 
+        FOLD_INFINITY_ARITHMETIC = enum.auto()
+
         @classmethod
         def all(self) -> ConstantFolding.Transformation:
             return functools.reduce(operator.or_, self.__members__.values())
@@ -239,3 +241,60 @@ def transform_fold_if(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
                 assert node.args[0].value == "False"
                 return node.args[2]
         return None
+
+    def transform_fold_infinity_arithmetic(self, node: ir.FunCall) -> Optional[ir.Node]:
+        if cpm.is_call_to(node, "plus"):
+            # `a + +/-inf` -> `+/-inf`
+            a, b = node.args
+            assert not (isinstance(a, ir.InfinityLiteral) and isinstance(b, ir.InfinityLiteral))
+            for arg in a, b:
+                if isinstance(arg, ir.InfinityLiteral):
+                    return arg
+
+        if cpm.is_call_to(node, "minimum"):
+            if ir.InfinityLiteral.NEGATIVE in node.args:
+                # `minimum(-inf, a)` -> `-inf`
+                return ir.InfinityLiteral.NEGATIVE
+            if ir.InfinityLiteral.POSITIVE in node.args:
+                # `minimum(inf, a)` -> `a`
+                a, b = node.args
+                return b if a == ir.InfinityLiteral.POSITIVE else a
+
+        if cpm.is_call_to(node, "maximum"):
+            if ir.InfinityLiteral.POSITIVE in node.args:
+                # `maximum(inf, a)` -> `inf`
+                return ir.InfinityLiteral.POSITIVE
+            if ir.InfinityLiteral.NEGATIVE in node.args:
+                # `maximum(-inf, a)` -> `a`
+                a, b = node.args
+                return b if a == ir.InfinityLiteral.NEGATIVE else a
+
+        if cpm.is_call_to(node, ("less", "less_equal")):
+            a, b = node.args
+            # we don't handle `inf < inf` or `-inf < -inf`.args
+            assert a != b or not isinstance(a, ir.InfinityLiteral)
+
+            # `-inf < v` -> `True`
+            # `v < inf` -> `True`
+            if a == ir.InfinityLiteral.NEGATIVE or b == ir.InfinityLiteral.POSITIVE:
+                return im.literal_from_value(True)
+            # `inf < v` -> `False`
+            # `v < -inf ` -> `False`
+            if a == ir.InfinityLiteral.POSITIVE or b == ir.InfinityLiteral.NEGATIVE:
+                return im.literal_from_value(False)
+
+        if cpm.is_call_to(node, ("greater", "greater_equal")):
+            a, b = node.args
+            # we don't handle `inf > inf` or `-inf > -inf`.args
+            assert a != b or not isinstance(a, ir.InfinityLiteral)
+
+            # `inf > v` -> `True`
+            # `v > -inf ` -> `True`
+            if a == ir.InfinityLiteral.POSITIVE or b == ir.InfinityLiteral.NEGATIVE:
+                return im.literal_from_value(True)
+            # `-inf > v` -> `False`
+            # `v > inf` -> `False`
+            if a == ir.InfinityLiteral.NEGATIVE or b == ir.InfinityLiteral.POSITIVE:
+                return im.literal_from_value(False)
+
+        return None

src/gt4py/next/iterator/transforms/cse.py

@@ -87,9 +87,19 @@ def _is_collectable_expr(node: itir.Node) -> bool:
         #  backend (single pass eager depth first visit approach)
         # do also not collect lifts or applied lifts as they become invisible to the lift inliner
         #  otherwise
-        if cpm.is_call_to(node, ("lift", "shift", "reduce", "map_")) or cpm.is_applied_lift(node):
+        # do also not collect index nodes because otherwise the right hand side of SetAts becomes a let statement
+        #  instead of an as_fieldop
+        if cpm.is_call_to(
+            node, ("lift", "shift", "reduce", "map_", "index")
+        ) or cpm.is_applied_lift(node):
             return False
         return True
+    # do also not collect make_tuple(index) nodes because otherwise the right hand side of SetAts becomes a let statement
+    #  instead of an as_fieldop
+    if cpm.is_call_to(node, "make_tuple") and all(
+        cpm.is_call_to(arg, "index") for arg in node.args
+    ):
+        return False
     elif isinstance(node, itir.Lambda):
         return True
 

src/gt4py/next/iterator/transforms/fuse_as_fieldop.py

@@ -450,7 +450,4 @@ def visit(self, node, **kwargs):
 
         node = super().visit(node, **kwargs)
 
-        if isinstance(node, itir.Expr) and hasattr(node.annex, "domain"):
-            node.annex.domain = node.annex.domain
-
         return node

src/gt4py/next/iterator/transforms/global_tmps.py

@@ -329,11 +329,16 @@ def create_global_tmps(
     This pass looks at all `as_fieldop` calls and transforms field-typed subexpressions of its
     arguments into temporaries.
     """
-    offset_provider_type = common.offset_provider_to_type(offset_provider)
+    # TODO(tehrengruber): document why to keep existing domains and add test
     program = infer_domain.infer_program(
-        program, offset_provider=offset_provider, symbolic_domain_sizes=symbolic_domain_sizes
+        program,
+        offset_provider=offset_provider,
+        symbolic_domain_sizes=symbolic_domain_sizes,
+        keep_existing_domains=True,
+    )
+    program = type_inference.infer(
+        program, offset_provider_type=common.offset_provider_to_type(offset_provider)
     )
-    program = type_inference.infer(program, offset_provider_type=offset_provider_type)
 
     if not uids:
         uids = eve_utils.UIDGenerator(prefix="__tmp")

src/gt4py/next/iterator/transforms/pass_manager.py

@@ -12,10 +12,12 @@
 from gt4py.next import common
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.transforms import (
+    concat_where,
     dead_code_elimination,
     fuse_as_fieldop,
     global_tmps,
     infer_domain,
+    infer_domain_ops,
     inline_dynamic_shifts,
     inline_fundefs,
     inline_lifts,
@@ -81,13 +83,19 @@ def apply_common_transforms(
     ir = inline_dynamic_shifts.InlineDynamicShifts.apply(
         ir
     )  # domain inference does not support dynamic offsets yet
+    ir = infer_domain_ops.InferDomainOps.apply(ir)
+    ir = concat_where.canonicalize_domain_argument(ir)
+
+    ir = concat_where.expand_tuple_args(ir, offset_provider_type=offset_provider_type)  # type: ignore[assignment]  # always an itir.Program
     ir = infer_domain.infer_program(
         ir,
         offset_provider=offset_provider,
         symbolic_domain_sizes=symbolic_domain_sizes,
     )
     ir = remove_broadcast.RemoveBroadcast.apply(ir)
 
+    ir = concat_where.transform_to_as_fieldop(ir)
+
     for _ in range(10):
         inlined = ir
 
@@ -183,6 +191,11 @@ def apply_fieldview_transforms(
     ir = inline_dynamic_shifts.InlineDynamicShifts.apply(
         ir
     )  # domain inference does not support dynamic offsets yet
+
+    ir = infer_domain_ops.InferDomainOps.apply(ir)
+    ir = concat_where.canonicalize_domain_argument(ir)
+    ir = ConstantFolding.apply(ir)  # type: ignore[assignment]  # always an itir.Program
+
     ir = infer_domain.infer_program(ir, offset_provider=offset_provider)
     ir = remove_broadcast.RemoveBroadcast.apply(ir)
     return ir