fix[next]: decide concat_where pruning from the concat_where's domain

src/gt4py/next/iterator/ir_utils/domain_utils.py

@@ -353,6 +353,20 @@ def promote_domain(
     return SymbolicDomain(domain.grid_type, dims_dict)
 
 
+def concat_where_branch_domain(
+    domain: SymbolicDomain, cond: SymbolicDomain, is_true_branch: bool
+) -> SymbolicDomain:
+    """
+    Return the part of `domain` on which one branch of a `concat_where` is selected.
+
+    Note: the complement is taken before promoting to the dimensions of `domain`, since
+    `domain_complement` requires each range to be infinite on exactly one side, which a
+    promoted dimension is not.
+    """
+    region = cond if is_true_branch else domain_complement(cond)
+    return domain_intersection(domain, promote_domain(region, domain.ranges.keys()))
+
+
 def is_finite(range_or_domain: SymbolicRange | SymbolicDomain) -> bool:
     """
     Return whether a range is unbounded in (at least) one direction.

src/gt4py/next/iterator/transforms/infer_domain.py

@@ -364,19 +364,14 @@ def _infer_concat_where(
     actual_domains: AccessedDomains = {}
     cond, true_field, false_field = expr.args
     symbolic_cond = domain_utils.SymbolicDomain.from_expr(cond)
-    cond_complement = domain_utils.domain_complement(symbolic_cond)
 
-    for arg in [true_field, false_field]:
+    for is_true_branch, arg in [(True, true_field), (False, false_field)]:
 
         @tree_map
         def mapper(d: NonTupleDomainAccess):
             if isinstance(d, DomainAccessDescriptor):
                 return d
-            promoted_cond = domain_utils.promote_domain(
-                symbolic_cond if arg == true_field else cond_complement,  # noqa: B023 # function is never used outside the loop
-                d.ranges.keys(),
-            )
-            return domain_utils.domain_intersection(d, promoted_cond)
+            return domain_utils.concat_where_branch_domain(d, symbolic_cond, is_true_branch)  # noqa: B023 # function is never used outside the loop
 
         domain_ = mapper(domain)
 

src/gt4py/next/iterator/transforms/prune_empty_concat_where.py

@@ -9,18 +9,19 @@
 from typing import TypeVar
 
 from gt4py.eve import NodeTranslator, PreserveLocationVisitor
-from gt4py.eve.extended_typing import Container, Self
-from gt4py.next import common
+from gt4py.eve.extended_typing import Self
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, domain_utils
 
 
-def _filter_domain(
-    domain: domain_utils.SymbolicDomain, dims: Container[common.Dimension]
-) -> domain_utils.SymbolicDomain:
-    return domain_utils.SymbolicDomain(
-        grid_type=domain.grid_type,
-        ranges={d: r for d, r in domain.ranges.items() if d in dims},
+def _covers(branch: itir.Expr, domain: domain_utils.SymbolicDomain) -> bool:
+    """Return whether `branch` can replace a `concat_where` accessed on `domain`."""
+    branch_domain = branch.annex.domain
+    # a branch domain is the selected region restricted to the dimensions of the branch, so
+    #  fewer dimensions mean the branch would have to be broadcast back before replacing
+    return (
+        isinstance(branch_domain, domain_utils.SymbolicDomain)
+        and branch_domain.ranges.keys() == domain.ranges.keys()
     )
 
 
@@ -32,14 +33,26 @@ class _PruneEmptyConcatWhere(PreserveLocationVisitor, NodeTranslator):
     """
     Prune `concat_where` expression with one branch never being accessed.
 
-    This pass requires domain inference to be executed before.
+    This pass requires domain inference to be executed before. In particular it relies on the
+     condition being in the canonical form that domain inference already requires, i.e. bounded
+     on exactly one side, as the complement of the condition is not defined otherwise.
 
     This pass the true and false branch values to be fields, not tuples of fields. Execute
      `gt4py.next.iterator.transforms.concat_where.expand_tuple_args` before.
 
-    >>> from gt4py.next.iterator.ir_utils import ir_makers as im
+    >>> from gt4py.next import common
+    >>> from gt4py.next.iterator.ir_utils import domain_utils, ir_makers as im
+    >>> from gt4py.next.iterator.transforms import infer_domain
     >>> IDim = common.Dimension("IDim")
-    >>> expr = im.concat_where(im.domain(common.GridType.UNSTRUCTURED, {IDim: (0, 0)}), "a", "b")
+    >>> expr = im.concat_where(
+    ...     im.domain(common.GridType.UNSTRUCTURED, {IDim: (10, itir.InfinityLiteral.POSITIVE)}),
+    ...     "a",
+    ...     "b",
+    ... )
+    >>> accessed = im.domain(common.GridType.UNSTRUCTURED, {IDim: (0, 10)})
+    >>> expr, _ = infer_domain.infer_expr(
+    ...     expr, domain_utils.SymbolicDomain.from_expr(accessed), offset_provider={}
+    ... )
     >>> assert prune_empty_concat_where(expr) == im.ref("b")
     """
 
@@ -60,18 +73,18 @@ def visit_FunCall(self, node: itir.FunCall) -> itir.Expr:
                 tb.annex.domain = node.annex.domain
                 return tb
 
+            domain = node.annex.domain
             cond = domain_utils.SymbolicDomain.from_expr(cond_expr)
-            if cond.empty():
-                return node.args[2]
-
-            tb_domain, fb_domain = (
-                _filter_domain(arg.annex.domain, cond.ranges.keys()) for arg in node.args[1:]
-            )
-            assert all(isinstance(d, domain_utils.SymbolicDomain) for d in (tb_domain, fb_domain))
-            if tb_domain.empty():
-                return node.args[2]
-            if fb_domain.empty():
-                return node.args[1]
+            if (
+                isinstance(domain, domain_utils.SymbolicDomain)
+                and cond.ranges.keys() <= domain.ranges.keys()
+            ):
+                for is_true_branch, other_branch in ((True, fb), (False, tb)):
+                    if domain_utils.concat_where_branch_domain(
+                        domain, cond, is_true_branch
+                    ).empty() and _covers(other_branch, domain):
+                        other_branch.annex.domain = domain
+                        return other_branch
 
         return node
 

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

@@ -407,3 +407,28 @@ def testee(
         out=out,
         ref=(ref0, ref1),
     )
+
+
+def test_concat_where_never_selected_branch_keeps_dimensions(cartesian_case):
+    # The bounds have to be literals: a branch can only be shown to be never selected
+    #  when the domain is known at compile time.
+    isize, ksize = 3, 4
+
+    @gtx.field_operator
+    def testee(interior: cases.IKField, profile: cases.KField) -> cases.IKField:
+        return concat_where(KDim < 0, interior, profile)
+
+    @gtx.program
+    def prog(interior: cases.IKField, profile: cases.KField, out: cases.IKField):
+        testee(interior, profile, out=out, domain={IDim: (0, 3), KDim: (0, 4)})
+
+    interior = cases.allocate(
+        cartesian_case, testee, "interior", domain=gtx.domain({IDim: isize, KDim: ksize})
+    )()
+    profile = cases.allocate(cartesian_case, testee, "profile", domain=gtx.domain({KDim: ksize}))()
+    out = cases.allocate(
+        cartesian_case, testee, cases.RETURN, domain=gtx.domain({IDim: isize, KDim: ksize})
+    )()
+
+    ref = np.broadcast_to(profile.asnumpy(), (isize, ksize))
+    cases.verify(cartesian_case, prog, interior, profile, out=out, ref=ref)

tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py

@@ -214,6 +214,30 @@ def test_promote_domain(testee_ranges, dimensions, expected_ranges):
         assert promoted == expected
 
 
+def test_concat_where_branch_domain_complements_before_promoting():
+    """Promoting first makes the complement undefined, so the order is not free.
+
+    `domain_complement` requires each range to be infinite on exactly one side, which
+    a promoted dimension is not.
+    """
+    domain = _make_domain({I: (0, 10), J: (0, 10)})
+    cond = domain_utils.SymbolicDomain(
+        grid_type=common.GridType.CARTESIAN,
+        ranges={J: domain_utils.SymbolicRange(5, itir.InfinityLiteral.POSITIVE)},
+    )
+
+    assert domain_utils.concat_where_branch_domain(domain, cond, True).as_expr() == im.domain(
+        common.GridType.CARTESIAN, {I: (0, 10), J: (5, 10)}
+    )
+    assert domain_utils.concat_where_branch_domain(domain, cond, False).as_expr() == im.domain(
+        common.GridType.CARTESIAN, {I: (0, 10), J: (0, 5)}
+    )
+
+    # the other order does not work
+    with pytest.raises(AssertionError):
+        domain_utils.domain_complement(domain_utils.promote_domain(cond, domain.ranges.keys()))
+
+
 def test_is_finite_symbolic_range():
     assert not domain_utils.is_finite(infinity_range)
     assert not domain_utils.is_finite(left_infinity_range)

tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py

@@ -15,10 +15,20 @@
 from gt4py.next.iterator.transforms.infer_domain import infer_expr
 from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas
 from gt4py.next.iterator.ir_utils import domain_utils
+from gt4py.next.type_system import type_specifications as ts
 
 Vertex = common.Dimension(value="Vertex", kind=common.DimensionKind.HORIZONTAL)
 K = common.Dimension(value="K", kind=common.DimensionKind.VERTICAL)
 
+float64 = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+vertex_k_field = ts.FieldType(dims=[Vertex, K], dtype=float64)
+vertex_field = ts.FieldType(dims=[Vertex], dtype=float64)
+k_field = ts.FieldType(dims=[K], dtype=float64)
+
+
+def _domain(ranges):
+    return domain_utils.SymbolicDomain.from_expr(im.domain(common.GridType.UNSTRUCTURED, ranges))
+
 
 @pytest.mark.parametrize(
     "accessed_domain, cond_domain, expected",
@@ -67,3 +77,46 @@ def test_prune_concat_where(accessed_domain, cond_domain, expected):
     actual = prune_empty_concat_where(testee)
     actual = InlineLambdas.apply(actual)
     assert actual == expected
+
+
+def _infer_concat_where(cond_range, true_branch_type, false_branch_type):
+    """A `concat_where` on `K` over the domain `Vertex: [0, 10), K: [0, 10)`."""
+    cond = im.domain(common.GridType.UNSTRUCTURED, {K: cond_range})
+    testee = im.concat_where(cond, im.ref("a", true_branch_type), im.ref("b", false_branch_type))
+    testee = canonicalize_domain_argument(testee)
+    testee, _ = infer_expr(
+        testee,
+        _domain({Vertex: (0, 10), K: (0, 10)}),
+        offset_provider={},
+    )
+    return testee
+
+
+@pytest.mark.parametrize(
+    "cond_range",
+    [
+        (itir.InfinityLiteral.NEGATIVE, 0),  # entirely below the domain
+        (10, itir.InfinityLiteral.POSITIVE),  # entirely above it
+    ],
+)
+def test_prune_branch_that_lacks_the_concat_dimension(cond_range):
+    """A never selected branch is pruned even when it does not have the concat dimension.
+
+    A branch's own domain is restricted to the dimensions of the branch's type, so for such
+    a branch it has no ranges at all and is vacuously non-empty. Whether a branch is selected
+    therefore has to be decided from the `concat_where`'s domain.
+    """
+    testee = _infer_concat_where(cond_range, vertex_field, vertex_k_field)
+
+    assert prune_empty_concat_where(testee) == im.ref("b", vertex_k_field)
+
+
+def test_no_prune_when_the_surviving_branch_lacks_a_dimension():
+    """Pruning must not silently drop a dimension.
+
+    The true branch is never selected, but the false branch does not have the `Vertex`
+    dimension, so replacing the `concat_where` by it would change the result.
+    """
+    testee = _infer_concat_where((itir.InfinityLiteral.NEGATIVE, 0), vertex_k_field, k_field)
+
+    assert prune_empty_concat_where(testee) == testee