From b13b96b40872111b5df49a5a2e2342eb9929216f Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 10 Aug 2026 15:40:17 +0200 Subject: [PATCH 1/2] fix[next]: decide concat_where pruning from the concat_where's domain `prune_empty_concat_where` decided whether a branch is ever selected from the branch's own domain, filtered to the dimensions of the condition. That is the wrong quantity in both directions. It is too weak: a branch domain is restricted to the dimensions of the branch's own type, so for a branch that does not have the concat dimension the filtered domain has no ranges at all and is never recognized as empty. A branch selected nowhere survives. It is also too strong: nothing checked that the surviving branch spans the same dimensions as the `concat_where`, so `concat_where(K < 0, a, b)` with `a: Field[[Vertex, K]]` and `b: Field[[K]]` was rewritten to `b`, replacing a two dimensional expression by a one dimensional one. Both follow from deriving the selected region from the `concat_where`'s own domain instead, which is what `infer_domain._infer_concat_where` already computes per branch. That computation moves to `domain_utils.concat_where_branch_domain` and both callers use it, so the requirement to take the complement before promoting -- `domain_complement` is undefined on a range that is infinite on both sides -- exists in one place. A branch is pruned only when its domain spans every dimension of the `concat_where`. Comparing domains rather than types also avoids depending on dimension order, and on a type that the pipeline does not set: `_infer_concat_where` rebuilds the node through `im.call`, which produces an untyped `FunCall`. --- .../next/iterator/ir_utils/domain_utils.py | 14 +++++ .../next/iterator/transforms/infer_domain.py | 9 +-- .../transforms/prune_empty_concat_where.py | 57 ++++++++++++------- .../ir_utils_tests/test_domain_utils.py | 24 ++++++++ .../test_prune_empty_concat_where.py | 53 +++++++++++++++++ 5 files changed, 128 insertions(+), 29 deletions(-) diff --git a/src/gt4py/next/iterator/ir_utils/domain_utils.py b/src/gt4py/next/iterator/ir_utils/domain_utils.py index 6bd28f6dbc..afa2a5e186 100644 --- a/src/gt4py/next/iterator/ir_utils/domain_utils.py +++ b/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. diff --git a/src/gt4py/next/iterator/transforms/infer_domain.py b/src/gt4py/next/iterator/transforms/infer_domain.py index c7e7d4f448..201c7315be 100644 --- a/src/gt4py/next/iterator/transforms/infer_domain.py +++ b/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) diff --git a/src/gt4py/next/iterator/transforms/prune_empty_concat_where.py b/src/gt4py/next/iterator/transforms/prune_empty_concat_where.py index 58b684acb2..946ce11c7a 100644 --- a/src/gt4py/next/iterator/transforms/prune_empty_concat_where.py +++ b/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 diff --git a/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py b/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py index 4aac031406..542eefe63e 100644 --- a/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py +++ b/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) diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py index c9c94096a5..0a0b708649 100644 --- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py +++ b/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 From 4084a075a7a692871bfb6f08f1360ff31e87aebc Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 13 Aug 2026 08:59:27 +0200 Subject: [PATCH 2/2] test[next]: cover a never selected concat_where branch end to end Every test in `test_concat_where.py` goes through `cases`, which produces `get_domain_range(out, ...)` bounds, so the domain is never known at compile time and `prune_empty_concat_where` can not fire in any of them. The new test uses a literal domain, which is what icon4py compiles with. Against the previous implementation it fails on every compiled backend, `gtfn`, `gtfn_imperative`, `dace_cpu` and `dace_cpu_noopt`, and passes on the embedded ones, which never lower. The surviving branch is substituted verbatim although it lacks a dimension of the `concat_where`, so the `SetAt` writes a one dimensional expression into a two dimensional target. With assertions disabled the mismatch reaches the backend and only the first line of the output is written. --- .../ffront_tests/test_concat_where.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py index e415bc7cfa..9063d850e5 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py +++ b/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)