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/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

@@ -443,6 +443,12 @@ 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)
+        true_type, false_type = node.args[1].type, node.args[2].type
+        if not isinstance(node.type, ts.TupleType) and any(
+            type_info.contains_local_field(t) for t in (true_type, false_type)
+        ):
+            true_branch = promote_to_list(true_type)(true_branch)
+            false_branch = promote_to_list(false_type)(false_branch)
         return im.concat_where(domain, true_branch, false_branch)
 
     def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:

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),

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py

@@ -8,13 +8,23 @@
 
 import numpy as np
 import pytest
-from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case
+from next_tests.integration_tests.cases import (
+    IDim,
+    JDim,
+    KDim,
+    V2E,
+    V2EDim,
+    Vertex,
+    cartesian_case,
+    unstructured_case,
+)
 from gt4py import next as gtx
-from gt4py.next import broadcast
+from gt4py.next import broadcast, common, neighbor_sum
 from gt4py.next.ffront.experimental import concat_where
 from next_tests.integration_tests import cases
 from next_tests.integration_tests.cases_utils import (
     exec_alloc_descriptor,
+    mesh_descriptor,
 )
 
 pytestmark = pytest.mark.uses_concat_where
@@ -390,6 +400,143 @@ def ref(interior0, boundary0, interior1, boundary1):
     cases.verify_with_default_data(cartesian_case, testee, ref)
 
 
+@pytest.mark.uses_tuple_returns
+def test_with_nested_tuples(cartesian_case, static_domains: bool):
+    @gtx.field_operator(static_domains=static_domains)
+    def testee(
+        interior0: cases.IJKField,
+        boundary0: cases.IJField,
+        interior1: cases.IJKField,
+        boundary1: cases.IJField,
+        interior2: cases.IJKField,
+        boundary2: cases.IJField,
+    ) -> tuple[cases.IJKField, tuple[cases.IJKField, cases.IJKField]]:
+        return concat_where(
+            KDim == 0,
+            (boundary0, (boundary1, boundary2)),
+            (interior0, (interior1, interior2)),
+        )
+
+    interiors = tuple(cases.allocate(cartesian_case, testee, f"interior{i}")() for i in range(3))
+    boundaries = tuple(cases.allocate(cartesian_case, testee, f"boundary{i}")() for i in range(3))
+    out = cases.allocate(cartesian_case, testee, cases.RETURN)()
+
+    k = np.arange(0, cartesian_case.default_sizes[KDim])
+    refs = tuple(
+        np.where(
+            k[np.newaxis, np.newaxis, :] == 0,
+            boundary.asnumpy()[:, :, np.newaxis],
+            interior.asnumpy(),
+        )
+        for boundary, interior in zip(boundaries, interiors)
+    )
+
+    cases.verify(
+        cartesian_case,
+        testee,
+        interiors[0],
+        boundaries[0],
+        interiors[1],
+        boundaries[1],
+        interiors[2],
+        boundaries[2],
+        out=out,
+        ref=(refs[0], (refs[1], refs[2])),
+    )
+
+
+@pytest.mark.uses_unstructured_shift
+@pytest.mark.uses_concat_where_with_list_output
+def test_with_local_field(unstructured_case, static_domains: bool):
+    @gtx.field_operator(static_domains=static_domains)
+    def testee(a: cases.EField, b: cases.EField) -> cases.VField:
+        t = concat_where(Vertex < 2, a(V2E), b(V2E))
+        return neighbor_sum(t, axis=V2EDim)
+
+    v2e_table = unstructured_case.offset_provider["V2E"].asnumpy()
+    vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2
+    cases.verify_with_default_data(
+        unstructured_case,
+        testee,
+        ref=lambda a, b: np.sum(
+            np.where(vertex_mask[:, np.newaxis], a[v2e_table], b[v2e_table]),
+            axis=1,
+            initial=0,
+            where=v2e_table != common._DEFAULT_SKIP_VALUE,
+        ),
+    )
+
+
+@pytest.mark.uses_tuple_returns
+@pytest.mark.uses_unstructured_shift
+@pytest.mark.uses_concat_where_with_list_output
+def test_with_tuples_of_local_fields(unstructured_case, static_domains: bool):
+    @gtx.field_operator(static_domains=static_domains)
+    def testee(
+        a: cases.EField,
+        b: cases.EField,
+        c: cases.EField,
+        d: cases.EField,
+    ) -> tuple[cases.VField, cases.VField]:
+        t = concat_where(Vertex < 2, (a(V2E), b(V2E)), (c(V2E), d(V2E)))
+        return neighbor_sum(t[0], axis=V2EDim), neighbor_sum(t[1], axis=V2EDim)
+
+    v2e_table = unstructured_case.offset_provider["V2E"].asnumpy()
+    vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2
+    cases.verify_with_default_data(
+        unstructured_case,
+        testee,
+        ref=lambda a, b, c, d: (
+            np.sum(
+                np.where(vertex_mask[:, np.newaxis], a[v2e_table], c[v2e_table]),
+                axis=1,
+                initial=0,
+                where=v2e_table != common._DEFAULT_SKIP_VALUE,
+            ),
+            np.sum(
+                np.where(vertex_mask[:, np.newaxis], b[v2e_table], d[v2e_table]),
+                axis=1,
+                initial=0,
+                where=v2e_table != common._DEFAULT_SKIP_VALUE,
+            ),
+        ),
+    )
+
+
+@pytest.mark.uses_tuple_returns
+@pytest.mark.uses_unstructured_shift
+@pytest.mark.uses_concat_where_with_list_output
+@pytest.mark.embedded_concat_where_infinite_domain
+def test_with_local_field_and_scalar(unstructured_case, static_domains: bool):
+    @gtx.field_operator(static_domains=static_domains)
+    def testee(a: cases.EField) -> tuple[cases.VField, cases.VField]:
+        return (
+            neighbor_sum(concat_where(Vertex < 2, a(V2E), 3), axis=V2EDim),
+            neighbor_sum(concat_where(Vertex < 2, 3, a(V2E)), axis=V2EDim),
+        )
+
+    v2e_table = unstructured_case.offset_provider["V2E"].asnumpy()
+    vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2
+    cases.verify_with_default_data(
+        unstructured_case,
+        testee,
+        ref=lambda a: (
+            np.sum(
+                np.where(vertex_mask[:, np.newaxis], a[v2e_table], 3),
+                axis=1,
+                initial=0,
+                where=v2e_table != common._DEFAULT_SKIP_VALUE,
+            ),
+            np.sum(
+                np.where(vertex_mask[:, np.newaxis], 3, a[v2e_table]),
+                axis=1,
+                initial=0,
+                where=v2e_table != common._DEFAULT_SKIP_VALUE,
+            ),
+        ),
+    )
+
+
 def test_nested_conditions_with_empty_branches(cartesian_case, static_domains: bool):
     @gtx.field_operator(static_domains=static_domains)
     def testee(interior: cases.IField, boundary: cases.IField, N: gtx.int32) -> cases.IField:

tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py

@@ -286,7 +286,10 @@ def conditional_wrong_arg_type(
 
     with pytest.raises(
         errors.DSLError,
-        match="Field arguments to 'where' must be of same dtype, got 'float32' != 'float64'.",
+        match=re.escape(
+            "Could not promote 'Field[[TDim], float32]' and 'Field[[TDim], float64]' "
+            "to common type in call to 'where'."
+        ),
     ):
         _ = FieldOperatorParser.apply_to_function(conditional_wrong_arg_type)
 

tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py

@@ -239,7 +239,7 @@ def domain_comparison(a: Field[[TDim], float], b: Field[[TDim], float]):
 
     with pytest.raises(
         errors.DSLError,
-        match="Field arguments to 'concat_where' must be of same dtype, got 'float64' != 'int32'.",
+        match="Could not promote 'float64' and 'int32' to common type in call to 'concat_where'.",
     ):
         _ = FieldOperatorParser.apply_to_function(domain_comparison)
 

tests/next_tests/unit_tests/type_system_tests/test_type_info.py

@@ -462,3 +462,17 @@ def test_return_type(
 )
 def test_needs_value_extraction(type_spec: ts.TypeSpec, expected: bool):
     assert type_info.needs_value_extraction(type_spec) is expected
+
+
+def test_promote_lists():
+    float64 = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+    V2EDim = Dimension("V2E", kind=DimensionKind.LOCAL)
+    C2EDim = Dimension("C2E", kind=DimensionKind.LOCAL)
+    const_list = ts.ListType(element_type=float64, offset_type=None)
+    v2e_list = ts.ListType(element_type=float64, offset_type=V2EDim)
+
+    assert type_info.promote(const_list, v2e_list) == v2e_list
+    with pytest.raises(ValueError, match="different offsets"):
+        type_info.promote(v2e_list, ts.ListType(element_type=float64, offset_type=C2EDim))
+    with pytest.raises(ValueError, match="non-lists"):
+        type_info.promote(v2e_list, float64)