Don't merge: ICON4Py staging

src/gt4py/next/common.py

@@ -109,6 +109,41 @@ def __add__(self, offset: int) -> Connectivity:
     def __sub__(self, offset: int) -> Connectivity:
         return self + (-offset)
 
+    def __gt__(self, value: core_defs.IntegralScalar) -> Domain:
+        return Domain(dims=(self,), ranges=(UnitRange(value + 1, Infinity.POSITIVE),))
+
+    def __ge__(self, value: core_defs.IntegralScalar) -> Domain:
+        return Domain(dims=(self,), ranges=(UnitRange(value, Infinity.POSITIVE),))
+
+    def __lt__(self, value: core_defs.IntegralScalar) -> Domain:
+        return Domain(dims=(self,), ranges=(UnitRange(Infinity.NEGATIVE, value),))
+
+    def __le__(self, value: core_defs.IntegralScalar) -> Domain:
+        # TODO add test
+        return Domain(dims=(self,), ranges=(UnitRange(Infinity.NEGATIVE, value + 1),))
+
+    def __eq__(self, value: Dimension | core_defs.IntegralScalar) -> bool | Domain:
+        if isinstance(value, Dimension):
+            return self.value == value.value
+        elif isinstance(value, core_defs.INTEGRAL_TYPES):
+            # TODO probably only within valid embedded context?
+            return Domain(dims=(self,), ranges=(UnitRange(value, value + 1),))
+        else:
+            return False
+
+    def __ne__(self, value: Dimension | core_defs.IntegralScalar) -> bool | tuple[Domain, Domain]:
+        # TODO add test
+        if isinstance(value, Dimension):
+            return self.value != value.value
+        elif isinstance(value, core_defs.INTEGRAL_TYPES):
+            # TODO probably only within valid embedded context?
+            return (
+                Domain(self, UnitRange(Infinity.NEGATIVE, value)),
+                Domain(self, UnitRange(value + 1, Infinity.POSITIVE)),
+            )
+        else:
+            return True
+
 
 class Infinity(enum.Enum):
     """Describes an unbounded `UnitRange`."""
@@ -500,6 +535,24 @@ def __and__(self, other: Domain) -> Domain:
         )
         return Domain(dims=broadcast_dims, ranges=intersected_ranges)
 
+    def __or__(self, other: Domain) -> Domain:
+        # TODO support arbitrary union of domains
+        # TODO add tests
+        if self.ndim > 1 or other.ndim > 1:
+            raise NotImplementedError("Union of multidimensional domains is not supported.")
+        if self.ndim == 0:
+            return other
+        if other.ndim == 0:
+            return self
+        sorted_ = sorted((self, other), key=lambda x: x.ranges[0].start)
+        if sorted_[0].ranges[0].stop >= sorted_[1].ranges[0].start:
+            return Domain(
+                dims=(self.dims[0],),
+                ranges=(UnitRange(sorted_[0].ranges[0].start, sorted_[1].ranges[0].stop),),
+            )
+        else:
+            return (sorted_[0], sorted_[1])
+
     @functools.cached_property
     def slice_at(self) -> utils.IndexerCallable[slice, Domain]:
         """

src/gt4py/next/embedded/nd_array_field.py

@@ -810,25 +810,6 @@ def _hyperslice(
 NdArrayField.register_builtin_func(fbuiltins.where, _make_builtin("where", "where"))
 
 
-def _compute_mask_slices(
-    mask: core_defs.NDArrayObject,
-) -> list[tuple[bool, slice]]:
-    """Take a 1-dimensional mask and return a sequence of mappings from boolean values to slices."""
-    # TODO: does it make sense to upgrade this naive algorithm to numpy?
-    assert mask.ndim == 1
-    cur = bool(mask[0].item())
-    ind = 0
-    res = []
-    for i in range(1, mask.shape[0]):
-        # Use `.item()` to extract the scalar from a 0-d array in case of e.g. cupy
-        if (mask_i := bool(mask[i].item())) != cur:
-            res.append((cur, slice(ind, i)))
-            cur = mask_i
-            ind = i
-    res.append((cur, slice(ind, mask.shape[0])))
-    return res
-
-
 def _trim_empty_domains(
     lst: Iterable[tuple[bool, common.Domain]],
 ) -> list[tuple[bool, common.Domain]]:
@@ -896,82 +877,108 @@ def _stack_domains(*domains: common.Domain, dim: common.Dimension) -> Optional[c
 
 def _concat(*fields: common.Field, dim: common.Dimension) -> common.Field:
     # TODO(havogt): this function could be extended to a general concat
-    # currently only concatenate along the given dimension and requires the fields to be ordered
+    # currently only concatenate along the given dimension
+    sorted_fields = sorted(fields, key=lambda f: f.domain[dim].unit_range.start)
 
     if (
-        len(fields) > 1
-        and not embedded_common.domain_intersection(*[f.domain for f in fields]).is_empty()
+        len(sorted_fields) > 1
+        and not embedded_common.domain_intersection(*[f.domain for f in sorted_fields]).is_empty()
     ):
         raise ValueError("Fields to concatenate must not overlap.")
-    new_domain = _stack_domains(*[f.domain for f in fields], dim=dim)
+    new_domain = _stack_domains(*[f.domain for f in sorted_fields], dim=dim)
     if new_domain is None:
         raise embedded_exceptions.NonContiguousDomain(f"Cannot concatenate fields along {dim}.")
-    nd_array_class = _get_nd_array_class(*fields)
+    nd_array_class = _get_nd_array_class(*sorted_fields)
     return nd_array_class.from_array(
         nd_array_class.array_ns.concatenate(
-            [nd_array_class.array_ns.broadcast_to(f.ndarray, f.domain.shape) for f in fields],
+            [
+                nd_array_class.array_ns.broadcast_to(f.ndarray, f.domain.shape)
+                for f in sorted_fields
+            ],
             axis=new_domain.dim_index(dim, allow_missing=False),
         ),
         domain=new_domain,
     )
 
 
+def _invert_domain(
+    domains: common.Domain | tuple[common.Domain],
+) -> common.Domain | tuple[common.Domain, ...]:
+    if not isinstance(domains, tuple):
+        domains = (domains,)
+
+    assert all(d.ndim == 1 for d in domains)
+    dim = domains[0].dims[0]
+    assert all(d.dims[0] == dim for d in domains)
+    sorted_domains = sorted(domains, key=lambda d: d.ranges[0].start)
+
+    result = []
+    if domains[0].ranges[0].start is not common.Infinity.NEGATIVE:
+        result.append(
+            common.Domain(
+                dims=(dim,),
+                ranges=(common.UnitRange(common.Infinity.NEGATIVE, domains[0].ranges[0].start),),
+            )
+        )
+    for i in range(len(sorted_domains) - 1):
+        if sorted_domains[i].ranges[0].stop != sorted_domains[i + 1].ranges[0].start:
+            result.append(
+                common.Domain(
+                    dims=(dim,),
+                    ranges=(
+                        common.UnitRange(
+                            sorted_domains[i].ranges[0].stop, sorted_domains[i + 1].ranges[0].start
+                        ),
+                    ),
+                )
+            )
+    if domains[-1].ranges[0].stop is not common.Infinity.POSITIVE:
+        result.append(
+            common.Domain(
+                dims=(dim,),
+                ranges=(common.UnitRange(domains[-1].ranges[0].stop, common.Infinity.POSITIVE),),
+            )
+        )
+    return tuple(result)
+
+
+def _intersect_multiple(
+    domain: common.Domain, domains: common.Domain | tuple[common.Domain]
+) -> tuple[common.Domain, ...]:
+    if not isinstance(domains, tuple):
+        domains = (domains,)
+
+    return tuple(
+        intersection
+        for d in domains
+        if not (intersection := embedded_common.domain_intersection(domain, d)).is_empty()
+    )
+
+
 def _concat_where(
-    mask_field: common.Field, true_field: common.Field, false_field: common.Field
+    masks: common.Domain | tuple[common.Domain, ...],
+    true_field: common.Field,
+    false_field: common.Field,
 ) -> common.Field:
-    cls_ = _get_nd_array_class(mask_field, true_field, false_field)
-    xp = cls_.array_ns
-    if mask_field.domain.ndim != 1:
+    if not isinstance(masks, tuple):
+        masks = (masks,)
+    if any(m.ndim for m in masks) != 1:
         raise NotImplementedError(
             "'concat_where': Can only concatenate fields with a 1-dimensional mask."
         )
-    mask_dim = mask_field.domain.dims[0]
+    mask_dim = masks[0].dims[0]
 
     # intersect the field in dimensions orthogonal to the mask, then all slices in the mask field have same domain
     t_broadcasted, f_broadcasted = _intersect_fields(true_field, false_field, ignore_dims=mask_dim)
 
-    # TODO(havogt): for clarity, most of it could be implemented on named_range in the masked dimension, but we currently lack the utils
-    # compute the consecutive ranges (first relative, then domain) of true and false values
-    mask_values_to_slices_mapping: Iterable[tuple[bool, slice]] = _compute_mask_slices(
-        mask_field.ndarray
-    )
-    mask_values_to_domain_mapping: Iterable[tuple[bool, common.Domain]] = (
-        (mask, mask_field.domain.slice_at[domain_slice])
-        for mask, domain_slice in mask_values_to_slices_mapping
-    )
-    # mask domains intersected with the respective fields
-    mask_values_to_intersected_domains_mapping: Iterable[tuple[bool, common.Domain]] = (
-        (
-            mask_value,
-            embedded_common.domain_intersection(
-                t_broadcasted.domain if mask_value else f_broadcasted.domain, mask_domain
-            ),
-        )
-        for mask_value, mask_domain in mask_values_to_domain_mapping
-    )
-
-    # remove the empty domains from the beginning and end
-    mask_values_to_intersected_domains_mapping = _trim_empty_domains(
-        mask_values_to_intersected_domains_mapping
-    )
-    if any(d.is_empty() for _, d in mask_values_to_intersected_domains_mapping):
-        raise embedded_exceptions.NonContiguousDomain(
-            f"In 'concat_where', cannot concatenate the following 'Domain's: {[d for _, d in mask_values_to_intersected_domains_mapping]}."
-        )
+    true_domains = _intersect_multiple(t_broadcasted.domain, masks)
+    t_slices = tuple(t_broadcasted[d] for d in true_domains)
 
-    # slice the fields with the domain ranges
-    transformed = [
-        t_broadcasted[d] if v else f_broadcasted[d]
-        for v, d in mask_values_to_intersected_domains_mapping
-    ]
+    inverted_masks = _invert_domain(masks)
+    false_domains = _intersect_multiple(f_broadcasted.domain, inverted_masks)
+    f_slices = tuple(f_broadcasted[d] for d in false_domains)
 
-    # stack the fields together
-    if transformed:
-        return _concat(*transformed, dim=mask_dim)
-    else:
-        result_domain = common.Domain(common.NamedRange(mask_dim, common.UnitRange(0, 0)))
-        result_array = xp.empty(result_domain.shape)
-    return cls_.from_array(result_array, domain=result_domain)
+    return _concat(*f_slices, *t_slices, dim=mask_dim)
 
 
 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

src/gt4py/next/ffront/foast_to_gtir.py

@@ -406,7 +406,12 @@ 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)
+        # TODO: use this case again. breaks domain inference in fused_velocity_advection_stencil_1_to_7
+        #  because some tuple elements are never accessed and the collapse tuple
+        #  does not propagate across concat where
+        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 +493,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,16 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     raise BackendNotSelectedError()
 
 
+@builtin_dispatch
+def concat_where(*args):
+    raise BackendNotSelectedError()
+
+
+@builtin_dispatch
+def in_(*args):
+    raise BackendNotSelectedError()
+
+
 UNARY_MATH_NUMBER_BUILTINS = {"abs", "neg"}
 UNARY_LOGICAL_BUILTINS = {"not_"}
 UNARY_MATH_FP_BUILTINS = {
@@ -494,6 +504,8 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     "scan",
     "tuple_get",
     "unstructured_domain",
+    "concat_where",
+    "in_",
     *ARITHMETIC_BUILTINS,
     *TYPE_BUILTINS,
 }

src/gt4py/next/iterator/embedded.py

@@ -1802,6 +1802,16 @@ 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.")
+
+
+@builtins.in_.register(EMBEDDED)
+def in_(*args):
+    raise NotImplementedError("To be implemented in frontend embedded.")
+
+
 def closure(
     domain_: runtime.CartesianDomain | runtime.UnstructuredDomain,
     sten: Callable[..., Any],

src/gt4py/next/iterator/ir.py

@@ -5,8 +5,10 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from __future__ import annotations
 
-from typing import ClassVar, List, Optional, Union
+import typing
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Union
 
 import gt4py.eve as eve
 from gt4py.eve import Coerced, SymbolName, SymbolRef
@@ -63,6 +65,22 @@ class NoneLiteral(Expr):
     _none_literal: int = 0
 
 
+class InfinityLiteral(Expr):
+    # TODO(tehrengruber): self referential `ClassVar` not supported in eve.
+    if TYPE_CHECKING:
+        POSITIVE: ClassVar[InfinityLiteral]
+        NEGATIVE: ClassVar[InfinityLiteral]
+
+    name: typing.Literal["POSITIVE", "NEGATIVE"]
+
+    def __str__(self):
+        return f"{type(self).__name__}.{self.name}"
+
+
+InfinityLiteral.NEGATIVE = InfinityLiteral(name="NEGATIVE")
+InfinityLiteral.POSITIVE = InfinityLiteral(name="POSITIVE")
+
+
 class OffsetLiteral(Expr):
     value: Union[int, str]
 
@@ -142,4 +160,5 @@ class Program(Node, ValidatedSymbolTableTrait):
 Program.__hash__ = Node.__hash__  # type: ignore[method-assign]
 SetAt.__hash__ = Node.__hash__  # type: ignore[method-assign]
 IfStmt.__hash__ = Node.__hash__  # type: ignore[method-assign]
+InfinityLiteral.__hash__ = Node.__hash__  # type: ignore[method-assign]
 Temporary.__hash__ = Node.__hash__  # type: ignore[method-assign]