feat[next]: Gtir concat where

pyproject.toml

@@ -327,7 +327,9 @@ markers = [
   'uses_unstructured_shift: tests that use a unstructured connectivity',
   'uses_max_over: tests that use the max_over builtin',
   'uses_mesh_with_skip_values: tests that use a mesh with skip values',
-  'checks_specific_error: tests that rely on the backend to produce a specific error message'
+  'checks_specific_error: tests that rely on the backend to produce a specific error message',
+  'uses_frontend_concat_where: tests that use the frontend concat_where builtin',
+  'uses_gtir_concat_where: tests that use the GTIR concat_where builtin'
 ]
 norecursedirs = ['dist', 'build', 'cpp_backend_tests/build*', '_local/*', '.*']
 testpaths = 'tests'

src/gt4py/eve/visitors.py

@@ -121,6 +121,21 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
         return None
 
 
+def _preserve_annex(
+    node: concepts.Node, new_node: concepts.Node, preserved_annex_attrs: tuple[str, ...]
+) -> None:
+    if preserved_annex_attrs and (old_annex := getattr(node, "__node_annex__", None)):
+        # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
+        new_annex_dict = new_node.annex.__dict__
+        for key in preserved_annex_attrs:
+            if (value := getattr(old_annex, key, NOTHING)) is not NOTHING:
+                # Note: The annex value of the new node might not be equal
+                # (in the sense that an equality comparison returns false),
+                # but in the context of the pass, they are equivalent.
+                # Therefore, we don't assert equality here.
+                new_annex_dict[key] = value
+
+
 class NodeTranslator(NodeVisitor):
     """Special `NodeVisitor` to translate nodes and trees.
 
@@ -158,13 +173,7 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
                     if (new_child := self.visit(child, **kwargs)) is not NOTHING
                 }
             )
-            if self.PRESERVED_ANNEX_ATTRS and (old_annex := getattr(node, "__node_annex__", None)):
-                # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
-                new_annex_dict = new_node.annex.__dict__
-                for key in self.PRESERVED_ANNEX_ATTRS:
-                    if (value := getattr(old_annex, key, NOTHING)) is not NOTHING:
-                        assert key not in new_annex_dict
-                        new_annex_dict[key] = value
+            _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
 
             return new_node
 
@@ -189,3 +198,11 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
             )
 
         return copy.deepcopy(node, memo=memo)
+
+    def visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
+        new_node = super().visit(node, **kwargs)
+
+        if isinstance(node, concepts.Node) and isinstance(new_node, concepts.Node):
+            _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
+
+        return new_node

src/gt4py/next/ffront/experimental.py

@@ -6,21 +6,31 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
-from typing import Tuple
+from typing import Tuple, TypeVar
 
 from gt4py._core import definitions as core_defs
 from gt4py.next import common
-from gt4py.next.ffront.fbuiltins import BuiltInFunction, FieldOffset, WhereBuiltinFunction
+from gt4py.next.ffront.fbuiltins import (
+    BuiltInFunction,
+    FieldOffset,
+    FieldT,
+    WhereLikeBuiltinFunction,
+)
 
 
 @BuiltInFunction
 def as_offset(offset_: FieldOffset, field: common.Field, /) -> common.Connectivity:
     raise NotImplementedError()
 
 
-@WhereBuiltinFunction
+_R = TypeVar("_R")
+DomainT = TypeVar("DomainT", bound=common.Field)
+ConcatWhereBuiltinFunction = WhereLikeBuiltinFunction[_R, DomainT, FieldT]
+
+
+@ConcatWhereBuiltinFunction
 def concat_where(
-    mask: common.Field,
+    mask: common.Domain,
     true_field: common.Field | core_defs.ScalarT | Tuple,
     false_field: common.Field | core_defs.ScalarT | Tuple,
     /,

src/gt4py/next/ffront/fbuiltins.py

@@ -66,6 +66,8 @@ def _type_conversion_helper(t: type) -> type[ts.TypeSpec] | tuple[type[ts.TypeSp
         return ts.OffsetType
     elif t is core_defs.ScalarT:
         return ts.ScalarType
+    elif t is common.Domain:
+        return ts.DomainType
     elif t is type:
         return (
             ts.FunctionType
@@ -135,14 +137,15 @@ def __gt_type__(self) -> ts.FunctionType:
         )
 
 
-MaskT = TypeVar("MaskT", bound=common.Field)
+MaskLikeT = TypeVar("MaskLikeT", bound=common.Field)
 FieldT = TypeVar("FieldT", bound=Union[common.Field, core_defs.Scalar, Tuple])
 
 
-class WhereBuiltinFunction(
-    BuiltInFunction[_R, [MaskT, FieldT, FieldT]], Generic[_R, MaskT, FieldT]
+class WhereLikeBuiltinFunction(
+    BuiltInFunction[_R, [MaskLikeT, FieldT, FieldT]],
+    Generic[_R, MaskLikeT, FieldT],
 ):
-    def __call__(self, mask: MaskT, true_field: FieldT, false_field: FieldT) -> _R:
+    def __call__(self, mask: MaskLikeT, true_field: FieldT, false_field: FieldT) -> _R:
         if isinstance(true_field, tuple) or isinstance(false_field, tuple):
             if not (isinstance(true_field, tuple) and isinstance(false_field, tuple)):
                 raise ValueError(
@@ -157,6 +160,10 @@ def __call__(self, mask: MaskT, true_field: FieldT, false_field: FieldT) -> _R:
         return super().__call__(mask, true_field, false_field)
 
 
+MaskT = TypeVar("MaskT", bound=common.Field)
+WhereBuiltinFunction = WhereLikeBuiltinFunction[_R, MaskT, FieldT]
+
+
 @BuiltInFunction
 def neighbor_sum(field: common.Field, /, axis: common.Dimension) -> common.Field:
     raise NotImplementedError()

src/gt4py/next/ffront/foast_passes/type_deduction.py

@@ -10,8 +10,8 @@
 
 import gt4py.next.ffront.field_operator_ast as foast
 from gt4py.eve import NodeTranslator, NodeVisitor, traits
-from gt4py.next import errors
-from gt4py.next.common import DimensionKind
+from gt4py.next import errors, utils
+from gt4py.next.common import DimensionKind, promote_dims
 from gt4py.next.ffront import (  # noqa
     dialect_ast_enums,
     experimental,
@@ -20,6 +20,7 @@
     type_specifications as ts_ffront,
 )
 from gt4py.next.ffront.foast_passes.utils import compute_assign_indices
+from gt4py.next.iterator import builtins
 from gt4py.next.type_system import type_info, type_specifications as ts, type_translation
 
 
@@ -566,16 +567,10 @@ def visit_Compare(self, node: foast.Compare, **kwargs: Any) -> foast.Compare:
             op=node.op, left=new_left, right=new_right, location=node.location, type=new_type
         )
 
-    def _deduce_compare_type(
+    def _deduce_arithmetic_compare_type(
         self, node: foast.Compare, *, left: foast.Expr, right: foast.Expr, **kwargs: Any
     ) -> Optional[ts.TypeSpec]:
-        # check both types compatible
-        for arg in (left, right):
-            if not type_info.is_arithmetic(arg.type):
-                raise errors.DSLError(
-                    arg.location, f"Type '{arg.type}' can not be used in operator '{node.op}'."
-                )
-
+        # e.g. `1 < 2`
         self._check_operand_dtypes_match(node, left=left, right=right)
 
         try:
@@ -592,6 +587,48 @@ def _deduce_compare_type(
                 f" in call to '{node.op}'.",
             ) from ex
 
+    def _deduce_dimension_compare_type(
+        self, node: foast.Compare, *, left: foast.Expr, right: foast.Expr, **kwargs: Any
+    ) -> Optional[ts.TypeSpec]:
+        # e.g. `IDim > 1`
+        index_type = ts.ScalarType(
+            kind=getattr(ts.ScalarKind, builtins.INTEGER_INDEX_BUILTIN.upper())
+        )
+
+        if isinstance(left.type, ts.DimensionType):
+            if not right.type == index_type:
+                raise errors.DSLError(
+                    right.location,
+                    f"Expected an {index_type}, but got '{right.type}' instead.",
+                )
+            return ts.DomainType(dims=[left.type.dim])
+        elif isinstance(right.type, ts.DimensionType):
+            if not left.type == index_type:
+                raise errors.DSLError(
+                    left.location,
+                    f"Expected an {index_type}, but got '{right.type}' instead.",
+                )
+            return ts.DomainType(dims=[right.type.dim])
+        else:
+            raise AssertionError()
+
+    def _deduce_compare_type(
+        self, node: foast.Compare, *, left: foast.Expr, right: foast.Expr, **kwargs: Any
+    ) -> Optional[ts.TypeSpec]:
+        # e.g. `1 < 1`
+        if all(type_info.is_arithmetic(arg) for arg in (left.type, right.type)):
+            return self._deduce_arithmetic_compare_type(node, left=left, right=right)
+        # e.g. `IDim > 1`
+        if any(isinstance(arg, ts.DimensionType) for arg in (left.type, right.type)):
+            return self._deduce_dimension_compare_type(node, left=left, right=right)
+
+        raise errors.DSLError(
+            left.location,
+            "Comparison operators can only be used between arithmetic types "
+            "(scalars, fields) or between a dimension and an index type "
+            "({builtins.INTEGER_INDEX_BUILTIN}).",
+        )
+
     def _deduce_binop_type(
         self, node: foast.BinOp, *, left: foast.Expr, right: foast.Expr, **kwargs: Any
     ) -> Optional[ts.TypeSpec]:
@@ -612,37 +649,48 @@ def _deduce_binop_type(
             dialect_ast_enums.BinaryOperator.BIT_OR,
             dialect_ast_enums.BinaryOperator.BIT_XOR,
         }
-        is_compatible = type_info.is_logical if node.op in logical_ops else type_info.is_arithmetic
-
-        # check both types compatible
-        for arg in (left, right):
-            if not is_compatible(arg.type):
-                raise errors.DSLError(
-                    arg.location, f"Type '{arg.type}' can not be used in operator '{node.op}'."
-                )
-
-        left_type = cast(ts.FieldType | ts.ScalarType, left.type)
-        right_type = cast(ts.FieldType | ts.ScalarType, right.type)
 
-        if node.op == dialect_ast_enums.BinaryOperator.POW:
-            return left_type
+        err_msg = f"Unsupported operand type(s) for {node.op}: '{left.type}' and '{right.type}'."
 
-        if node.op == dialect_ast_enums.BinaryOperator.MOD and not type_info.is_integral(
-            right_type
+        if isinstance(left.type, (ts.ScalarType, ts.FieldType)) and isinstance(
+            right.type, (ts.ScalarType, ts.FieldType)
         ):
-            raise errors.DSLError(
-                arg.location,
-                f"Type '{right_type}' can not be used in operator '{node.op}', it only accepts 'int'.",
+            is_compatible = (
+                type_info.is_logical if node.op in logical_ops else type_info.is_arithmetic
             )
+            for arg in (left, right):
+                if not is_compatible(arg.type):
+                    raise errors.DSLError(arg.location, err_msg)
 
-        try:
-            return type_info.promote(left_type, right_type)
-        except ValueError as ex:
-            raise errors.DSLError(
-                node.location,
-                f"Could not promote '{left_type}' and '{right_type}' to common type"
-                f" in call to '{node.op}'.",
-            ) from ex
+            if node.op == dialect_ast_enums.BinaryOperator.POW:
+                return left.type
+
+            if node.op == dialect_ast_enums.BinaryOperator.MOD and not type_info.is_integral(
+                right.type
+            ):
+                raise errors.DSLError(
+                    arg.location,
+                    f"Type '{right.type}' can not be used in operator '{node.op}', it only accepts 'int'.",
+                )
+
+            try:
+                return type_info.promote(left.type, right.type)
+            except ValueError as ex:
+                raise errors.DSLError(
+                    node.location,
+                    f"Could not promote '{left.type}' and '{right.type}' to common type"
+                    f" in call to '{node.op}'.",
+                ) from ex
+        elif isinstance(left.type, ts.DomainType) and isinstance(right.type, ts.DomainType):
+            if node.op not in logical_ops:
+                raise errors.DSLError(
+                    node.location,
+                    f"{err_msg} Operator "
+                    f"must be one of {', '.join((str(op) for op in logical_ops))}.",
+                )
+            return ts.DomainType(dims=promote_dims(left.type.dims, right.type.dims))
+        else:
+            raise errors.DSLError(node.location, err_msg)
 
     def _check_operand_dtypes_match(
         self, node: foast.BinOp | foast.Compare, left: foast.Expr, right: foast.Expr
@@ -908,6 +956,7 @@ def _visit_where(self, node: foast.Call, **kwargs: Any) -> foast.Call:
             )
 
         try:
+            # TODO(tehrengruber): the construct_tuple_type function doesn't look correct
             if isinstance(true_branch_type, ts.TupleType) and isinstance(
                 false_branch_type, ts.TupleType
             ):
@@ -943,7 +992,39 @@ def _visit_where(self, node: foast.Call, **kwargs: Any) -> foast.Call:
             location=node.location,
         )
 
-    _visit_concat_where = _visit_where
+    def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> foast.Call:
+        mask_type, true_branch_type, false_branch_type = (arg.type for arg in node.args)
+
+        assert isinstance(mask_type, ts.DomainType)
+        assert all(
+            isinstance(el, (ts.FieldType, ts.ScalarType))
+            for arg in (true_branch_type, false_branch_type)
+            for el in type_info.primitive_constituents(arg)
+        )
+
+        @utils.tree_map(
+            collection_type=ts.TupleType,
+            result_collection_constructor=lambda el: ts.TupleType(types=list(el)),
+        )
+        def deduce_return_type(tb: ts.FieldType | ts.ScalarType, fb: ts.FieldType | ts.ScalarType):
+            if (t_dtype := type_info.extract_dtype(tb)) != (f_dtype := type_info.extract_dtype(fb)):
+                raise errors.DSLError(
+                    node.location,
+                    f"Field arguments must be of same dtype, got '{t_dtype}' != '{f_dtype}'.",
+                )
+            return_dims = promote_dims(mask_type.dims, type_info.promote(tb, fb).dims)
+            return_type = ts.FieldType(dims=return_dims, dtype=type_info.promote(t_dtype, f_dtype))
+            return return_type
+
+        return_type = deduce_return_type(true_branch_type, false_branch_type)
+
+        return foast.Call(
+            func=node.func,
+            args=node.args,
+            kwargs=node.kwargs,
+            type=return_type,
+            location=node.location,
+        )
 
     def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> foast.Call:
         arg_type = cast(ts.FieldType | ts.ScalarType, node.args[0].type)

src/gt4py/next/ffront/foast_to_gtir.py

@@ -224,9 +224,16 @@ def visit_Assign(
     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):
+        if isinstance(node.type, ts.DimensionType):
+            return itir.AxisLiteral(value=node.type.dim.value, kind=node.type.dim.kind)
+        raise AssertionError()
+
     def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> itir.Expr:
         return im.tuple_get(node.index, self.visit(node.value, **kwargs))
 
@@ -394,7 +401,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)
+        return im.concat_where(domain, true_branch, false_branch)
 
     def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         expr = self.visit(node.args[0], **kwargs)
@@ -477,7 +486,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)
     ):