fix[next]: concat_where crashes on neighbor-list fields

pyproject.toml

@@ -310,6 +310,7 @@ markers = [
   '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_concat_where_with_list_output: tests that use concat_where with a neighbor-list result',
   'embedded_concat_where_infinite_domain: tests with concat_where resulting in an infinite domain',
   'embedded_concat_where_non_contiguous_domain: tests with concat_where on non-contiguous domains',
   'uses_program_metrics: tests that require backend support for program metrics',

src/gt4py/next/ffront/fbuiltins.py

@@ -243,6 +243,10 @@ def __call__(
     ) -> Tuple: ...
 
     def __call__(self, cond: CondT, true_field: FieldT1, false_field: FieldT2) -> _R:  # type: ignore[misc] # supposedly this signature does not accept all the possible args allowed by the overloads ??
+        if isinstance(true_field, named_collections.CUSTOM_NAMED_COLLECTION_TYPES):
+            return named_collections.tree_map_named_collection(lambda t, f: self(cond, t, f))(  # type: ignore[return-value] # `NamedCollection` is not `_R`
+                true_field, false_field
+            )
         if isinstance(true_field, tuple) or isinstance(false_field, tuple):
             if not (isinstance(true_field, tuple) and isinstance(false_field, tuple)):
                 raise ValueError(

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

@@ -871,8 +871,12 @@ def _visit_math_built_in(self, node: foast.Call, **kwargs: Any) -> foast.Call:
             )
         elif func_name in fbuiltins.BINARY_MATH_NUMBER_BUILTIN_NAMES:
             try:
-                return_type = type_info.promote(
-                    *((cast(ts.FieldType | ts.ScalarType, arg.type)) for arg in node.args)
+                return_type = cast(
+                    # a `ListType` only occurs at the ITIR level, never in the frontend
+                    ts.FieldType | ts.ScalarType,
+                    type_info.promote(
+                        *((cast(ts.FieldType | ts.ScalarType, arg.type)) for arg in node.args)
+                    ),
                 )
             except ValueError as ex:
                 raise errors.DSLError(node.location, error_msg_preamble) from ex
@@ -1020,15 +1024,17 @@ def _deduce_where_return_type(
         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)):
+            try:
+                promoted = type_info.promote(tb, fb)
+            except ValueError as ex:
                 raise errors.DSLError(
                     location,
-                    f"Field arguments to '{func_name}' must be of same dtype, got '{t_dtype}' != "
-                    f"'{f_dtype}'.",
-                )
-            return_dims = promote_dims(cond_dims, type_info.extract_dims(type_info.promote(tb, fb)))
-            return_type = ts.FieldType(dims=return_dims, dtype=t_dtype)
-            return return_type
+                    f"Could not promote '{tb}' and '{fb}' to common type in call to '{func_name}'.",
+                ) from ex
+            return ts.FieldType(
+                dims=promote_dims(cond_dims, type_info.extract_dims(promoted)),
+                dtype=type_info.extract_dtype(promoted),
+            )
 
         return deduce_return_type(true_branch, false_branch)
 

src/gt4py/next/ffront/foast_to_gtir.py

@@ -415,8 +415,8 @@ def _visit_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         # TODO(tehrengruber): For tuples we expand the tuple structure via `process_elements`
         #  instead of emitting `tree_map_tuple` so mixed field types are supported,
         #  e.g. (local field, regular field).
-        if not isinstance(node.type, ts.TupleType):  # to keep the IR simpler
-            return self._lower_and_map("if_", *node.args)
+        if not isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)):
+            return self._lower_and_map("if_", *node.args)  # to keep the IR simpler
 
         cond_ = self.visit(node.args[0])
         cond_symref_name = f"__cond_{itir.lenient_ir_fingerprinter(cond_)}"
@@ -443,7 +443,20 @@ def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         # TODO(tehrengruber): Use `tree_map_tuple` when the domain inference is able to handle
         #  lambda functions (with the results domain depending on the caller / args)
         domain, true_branch, false_branch = self.visit(node.args, **kwargs)
-        return im.concat_where(domain, true_branch, false_branch)
+
+        def create_concat_where(
+            true_: itir.Expr, false_: itir.Expr, arg_types: tuple[ts.TypeSpec, ts.TypeSpec]
+        ) -> itir.FunCall:
+            if any(type_info.contains_local_field(t) for t in arg_types):
+                true_, false_ = (promote_to_list(t)(e) for t, e in zip(arg_types, (true_, false_)))
+            return im.concat_where(domain, true_, false_)
+
+        branch_types = (node.args[1].type, node.args[2].type)
+        if not isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)):
+            return create_concat_where(true_branch, false_branch, branch_types)  # to keep the IR simpler
+        return lowering_utils.process_elements(
+            create_concat_where, (true_branch, false_branch), node.type, arg_types=branch_types
+        )
 
     def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         return im.call("broadcast")(*self.visit(node.args, **kwargs))

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

@@ -216,6 +216,8 @@ def if_(
     #  want this, but for roundtrip it is totally fine.
     # assert true_branch == false_branch  # noqa: ERA001
 
+    if isinstance(true_branch, ts.ListType) and isinstance(false_branch, ts.ListType):
+        return type_info.promote(true_branch, false_branch)
     return true_branch
 
 
@@ -279,18 +281,11 @@ def deduce_return_type(tb: ts.FieldType | ts.ScalarType, fb: ts.FieldType | ts.S
         if any(isinstance(b, ts.DeferredType) for b in [tb, fb]):
             return ts.DeferredType(constraint=ts.FieldType)
 
-        tb_dtype, fb_dtype = (type_info.extract_dtype(b) for b in [tb, fb])
-
-        assert tb_dtype == fb_dtype, (
-            f"Field arguments to 'concat_where' must be of same dtype, got '{tb_dtype}' != '{fb_dtype}'."
-        )
-        dtype = tb_dtype
-
-        return_dims = common.promote_dims(
-            domain.dims, type_info.extract_dims(type_info.promote(tb, fb))
+        promoted = type_info.promote(tb, fb)
+        return ts.FieldType(
+            dims=common.promote_dims(domain.dims, type_info.extract_dims(promoted)),
+            dtype=type_info.extract_dtype(promoted),
         )
-        return_type = ts.FieldType(dims=return_dims, dtype=dtype)
-        return return_type
 
     return deduce_return_type(true_field, false_field)
 

src/gt4py/next/type_system/type_info.py

@@ -9,7 +9,7 @@
 import functools
 import types
 from collections.abc import Callable, Collection, Iterable, Iterator
-from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, cast, overload
+from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, overload
 
 import numpy as np
 
@@ -556,15 +556,18 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool:
 
 
 def promote(
-    *types: ts.FieldType | ts.ScalarType, always_field: bool = False
-) -> ts.FieldType | ts.ScalarType:
+    *types: ts.FieldType | ts.ScalarType | ts.ListType, always_field: bool = False
+) -> ts.FieldType | ts.ScalarType | ts.ListType:
     """
-    Promote a set of field or scalar types to a common type.
+    Promote a set of field, scalar or list types to a common type.
 
     The resulting type is defined on all dimensions of the arguments, respecting
     the individual order of the dimensions of each argument (see
     :func:`common.promote_dims` for more details).
 
+    `ListType`s promote only with each other: their element types must be equal, and an
+    `offset_type` of `None` (a list from `make_const_list`) is compatible with any other.
+
     >>> dtype = ts.ScalarType(kind=ts.ScalarKind.INT64)
     >>> I, J, K = (common.Dimension(value=dim) for dim in ["I", "J", "K"])
     >>> promoted: ts.FieldType = promote(
@@ -578,18 +581,45 @@ def promote(
     ... )
     >>> promoted.dims == [I, J, K] and promoted.dtype == dtype
     True
+
+    >>> V2E = common.Dimension(value="V2E", kind=common.DimensionKind.LOCAL)
+    >>> list_dtype = ts.ListType(element_type=dtype, offset_type=V2E)
+    >>> promote(
+    ...     ts.FieldType(dims=[I], dtype=list_dtype),
+    ...     ts.FieldType(dims=[I, J], dtype=list_dtype),
+    ... ) == ts.FieldType(dims=[I, J], dtype=list_dtype)
+    True
+
+    >>> const_list_dtype = ts.ListType(element_type=dtype, offset_type=None)
+    >>> promote(
+    ...     ts.FieldType(dims=[I], dtype=const_list_dtype),
+    ...     ts.FieldType(dims=[I], dtype=list_dtype),
+    ... ) == ts.FieldType(dims=[I], dtype=list_dtype)
+    True
     """
     if not always_field and all(isinstance(type_, ts.ScalarType) for type_ in types):
         if not all(type_ == types[0] for type_ in types):
             raise ValueError("Could not promote scalars of different dtype (not implemented).")
-        if not all(type_.shape is None for type_ in types):  # type: ignore[union-attr]
+        if isinstance(types[0], ts.ScalarType) and types[0].shape is not None:
             raise NotImplementedError("Shape promotion not implemented.")
         return types[0]
+    elif not always_field and any(isinstance(type_, ts.ListType) for type_ in types):
+        lists = [type_ for type_ in types if isinstance(type_, ts.ListType)]
+        if len(lists) != len(types):
+            raise ValueError("Could not promote lists together with non-lists.")
+        if not all(list_.element_type == lists[0].element_type for list_ in lists):
+            raise ValueError("Could not promote lists of different element type (not implemented).")
+        offset_types = [list_.offset_type for list_ in lists if list_.offset_type is not None]
+        if not all(offset_type == offset_types[0] for offset_type in offset_types):
+            raise ValueError("Could not promote lists over different offsets.")
+        return ts.ListType(
+            element_type=lists[0].element_type,
+            offset_type=offset_types[0] if offset_types else None,
+        )
     elif all(isinstance(type_, (ts.ScalarType, ts.FieldType)) for type_ in types):
         dims = common.promote_dims(*(extract_dims(type_) for type_ in types))
-        extracted_dtypes = [extract_dtype(type_) for type_ in types]
-        assert all(isinstance(dtype, ts.ScalarType) for dtype in extracted_dtypes)
-        dtype = cast(ts.ScalarType, promote(*extracted_dtypes))  # type: ignore[arg-type] # checked is `ScalarType`
+        dtype = promote(*(extract_dtype(type_) for type_ in types))
+        assert isinstance(dtype, (ts.ScalarType, ts.ListType))
 
         return ts.FieldType(dims=dims, dtype=dtype)
     raise TypeError("Expected a 'FieldType' or 'ScalarType'.")

tests/next_tests/definitions.py

@@ -125,6 +125,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum):
 USES_PROGRAM_METRICS = "uses_program_metrics"
 USES_SCALAR_IN_DOMAIN_AND_FO = "uses_scalar_in_domain_and_fo"
 USES_CONCAT_WHERE = "uses_concat_where"
+USES_CONCAT_WHERE_WITH_LIST_OUTPUT = "uses_concat_where_with_list_output"
 EMBEDDED_CONCAT_WHERE_INFINITE_DOMAIN = "embedded_concat_where_infinite_domain"
 EMBEDDED_CONCAT_WHERE_NON_CONTIGUOUS_DOMAIN = "embedded_concat_where_non_contiguous_domain"
 USES_PROGRAM_WITH_SLICED_OUT_ARGUMENTS = "uses_program_with_sliced_out_arguments"
@@ -166,6 +167,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum):
     + [
         (USES_CAN_DEREF, XFAIL, UNSUPPORTED_MESSAGE),
         (USES_COMPOSITE_SHIFTS, XFAIL, UNSUPPORTED_MESSAGE),
+        (USES_CONCAT_WHERE_WITH_LIST_OUTPUT, XFAIL, UNSUPPORTED_MESSAGE),
         (USES_LIFT, XFAIL, UNSUPPORTED_MESSAGE),
         (USES_REDUCE_WITH_LAMBDA, XFAIL, UNSUPPORTED_MESSAGE),
         (USES_SCAN_IN_STENCIL, XFAIL, BINDINGS_UNSUPPORTED_MESSAGE),