feat[next]: GTIR concat_where frontend (remove chained comparison) (GridTools#1998)

GTIR concat_where frontend parts extracted from GridTools#1713. Credits to @SF-N
and @tehrengruber .

This PR removes the support for chained comparisons like `field_a <
field_b < field_c` which cannot be supported in embedded (same as `0 <
Dim < 42`) because it translates to a boolean scalar comparison as it is
evaluated as `(field_a < field_b) and (field_b < field_c)` with `and`
instead of `&`.

---------

Co-authored-by: tehrengruber <[email protected]>

Author: 2 people

pyproject.toml

@@ -302,6 +302,7 @@ 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',
+  'uses_concat_where: tests that use the concat_where builtin',
   'uses_program_metrics: tests that require backend support for program metrics',
   'checks_specific_error: tests that rely on the backend to produce a specific error message'
 ]

src/gt4py/next/embedded/nd_array_field.py

@@ -974,7 +974,7 @@ def _concat_where(
     return cls_.from_array(result_array, domain=result_domain)
 
 
-NdArrayField.register_builtin_func(experimental.concat_where, _concat_where)  # type: ignore[arg-type]
+NdArrayField.register_builtin_func(experimental.concat_where, _concat_where)  # type: ignore[arg-type] # TODO(havogt): this is still the "old" concat_where, needs to be replaced in a next PR
 
 
 def _make_reduction(

src/gt4py/next/ffront/ast_passes/unchain_compares.py

@@ -46,7 +46,7 @@ def visit_Compare(self, node: ast.Compare) -> ast.Compare | ast.BinOp:
 
         # the remainder of the chain -> right branch of the new tree
         # example: ``b > c > d``
-        remaining_chain = copy.copy(node)
+        remaining_chain = copy.deepcopy(node)
         remaining_chain.left = remaining_chain.comparators.pop(0)
         remaining_chain.ops.pop(0)
 

src/gt4py/next/ffront/experimental.py

@@ -20,7 +20,7 @@ def as_offset(offset_: FieldOffset, field: common.Field, /) -> common.Connectivi
 
 @WhereBuiltinFunction
 def concat_where(
-    mask: common.Field,
+    cond: 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,14 @@ def __gt_type__(self) -> ts.FunctionType:
         )
 
 
-MaskT = TypeVar("MaskT", bound=common.Field)
+CondT = TypeVar("CondT", bound=Union[common.Field, common.Domain])
 FieldT = TypeVar("FieldT", bound=Union[common.Field, core_defs.Scalar, Tuple])
 
 
 class WhereBuiltinFunction(
-    BuiltInFunction[_R, [MaskT, FieldT, FieldT]], Generic[_R, MaskT, FieldT]
+    BuiltInFunction[_R, [CondT, FieldT, FieldT]], Generic[_R, CondT, FieldT]
 ):
-    def __call__(self, mask: MaskT, true_field: FieldT, false_field: FieldT) -> _R:
+    def __call__(self, cond: CondT, 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(
@@ -153,8 +155,8 @@ def __call__(self, mask: MaskT, true_field: FieldT, false_field: FieldT) -> _R:
                 raise ValueError(
                     "Tuple of different size not allowed."
                 )  # TODO(havogt) find a strategy to unify parsing and embedded error messages
-            return tuple(self(mask, t, f) for t, f in zip(true_field, false_field))  # type: ignore[return-value] # `tuple` is not `_R`
-        return super().__call__(mask, true_field, false_field)
+            return tuple(self(cond, t, f) for t, f in zip(true_field, false_field))  # type: ignore[return-value] # `tuple` is not `_R`
+        return super().__call__(cond, true_field, false_field)
 
 
 @BuiltInFunction

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
 
 
@@ -428,7 +429,7 @@ def visit_IfStmt(self, node: foast.IfStmt, **kwargs: Any) -> foast.IfStmt:
         if not isinstance(new_node.condition.type, ts.ScalarType):
             raise errors.DSLError(
                 node.location,
-                "Condition for 'if' must be scalar, " f"got '{new_node.condition.type}' instead.",
+                f"Condition for 'if' must be scalar, got '{new_node.condition.type}' instead.",
             )
 
         if new_node.condition.type.kind != ts.ScalarKind.BOOL:
@@ -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,51 @@ 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())
+        )
+
+        def error_msg(left: ts.TypeSpec, right: ts.TypeSpec) -> str:
+            return f"Dimension comparison needs to be between a 'Dimension' and index of type '{index_type}', got '{left}' and '{right}'."
+
+        if isinstance(left.type, ts.DimensionType):
+            if not right.type == index_type:
+                raise errors.DSLError(
+                    right.location,
+                    error_msg(left.type, right.type),
+                )
+            return ts.DomainType(dims=[left.type.dim])
+        elif isinstance(right.type, ts.DimensionType):
+            if not left.type == index_type:
+                raise errors.DSLError(
+                    left.location,
+                    error_msg(left.type, right.type),
+                )
+            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 +652,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 +959,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 +995,43 @@ 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:
+        cond_type, true_branch_type, false_branch_type = (arg.type for arg in node.args)
+
+        assert isinstance(cond_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
+        ) -> ts.FieldType:
+            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(
+                cond_type.dims, type_info.extract_dims(type_info.promote(tb, fb))
+            )
+            return_type = ts.FieldType(dims=return_dims, dtype=t_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/func_to_foast.py

@@ -10,6 +10,7 @@
 
 import ast
 import builtins
+import textwrap
 import typing
 from typing import Any, Callable, Iterable, Mapping, Type
 
@@ -144,12 +145,11 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
     """
 
     @classmethod
-    def _preprocess_definition_ast(cls, definition_ast: ast.AST) -> ast.AST:
-        sta = StringifyAnnotationsPass.apply(definition_ast)
-        ssa = SingleStaticAssignPass.apply(sta)
-        sat = SingleAssignTargetPass.apply(ssa)
-        ucc = UnchainComparesPass.apply(sat)
-        return ucc
+    def _preprocess_definition_ast(cls, ast: ast.AST) -> ast.AST:
+        ast = StringifyAnnotationsPass.apply(ast)
+        ast = SingleStaticAssignPass.apply(ast)
+        ast = SingleAssignTargetPass.apply(ast)
+        return ast
 
     @classmethod
     def _postprocess_dialect_ast(
@@ -474,10 +474,20 @@ def _visit_stmts(
 
     def visit_Compare(self, node: ast.Compare, **kwargs: Any) -> foast.Compare:
         loc = self.get_location(node)
+
         if len(node.ops) != 1 or len(node.comparators) != 1:
-            # Remove comparison chains in a preprocessing pass
-            # TODO: maybe add a note to the error about preprocessing passes?
-            raise errors.UnsupportedPythonFeatureError(loc, "comparison chains")
+            refactored = UnchainComparesPass.apply(node)
+            raise errors.DSLError(
+                loc,
+                textwrap.dedent(
+                    f"""
+                    Comparison chains are not allowed. Please replace
+                        {ast.unparse(node)}
+                    by
+                        {ast.unparse(refactored)}
+                    """,
+                ),
+            )
         return foast.Compare(
             op=self.visit(node.ops[0]),
             left=self.visit(node.left),

src/gt4py/next/iterator/type_system/inference.py

@@ -455,7 +455,7 @@ def visit_Program(self, node: itir.Program, *, ctx) -> it_ts.ProgramType:
 
     def visit_Temporary(self, node: itir.Temporary, *, ctx) -> ts.FieldType | ts.TupleType:
         domain = self.visit(node.domain, ctx=ctx)
-        assert isinstance(domain, it_ts.DomainType)
+        assert isinstance(domain, ts.DomainType)
         assert domain.dims != "unknown"
         assert node.dtype
         return type_info.apply_to_primitive_constituents(

src/gt4py/next/iterator/type_system/type_specifications.py

@@ -16,10 +16,6 @@ class NamedRangeType(ts.TypeSpec):
     dim: common.Dimension
 
 
-class DomainType(ts.DataType):
-    dims: list[common.Dimension]
-
-
 class OffsetLiteralType(ts.TypeSpec):
     value: ts.ScalarType | str