fix[next]: Improve concat_where pruning

[tehrengruber]

Problem

A concat_where implicitly broadcasts each branch to the dimensions the branch does not have itself: in concat_where(K < 0, a, b) with a: Field[[Vertex]] and b: Field[[Vertex, K]], accessed on Vertex: [0, 10), K: [0, 10), the true branch is effectively broadcast(a, (Vertex, K)). This broadcast exists only implicitly — nothing in the IR represents it, and consequently the inferred domain of a branch is restricted to the dimensions of the branch's own type: the domain of the true_branch argument a is just Vertex: [0, 10). prune_empty_concat_where decided whether a branch is ever selected from exactly these branch domains (filtered to the dimensions of the condition), which fails in two ways:

• Never-selected branches survived. In the example, a is selected where K < 0, i.e. nowhere within K: [0, 10) — but the empty range K: [0, 0) proving this is exactly the one missing from a's domain, so a was never pruned.
• Pruning dropped dimensions. With the condition flipped, concat_where(0 <= K, a, b) selects a everywhere and was pruned to plain a — replacing the two-dimensional expression by a one-dimensional one, as nothing accounted for the surviving branch relying on the implicit broadcast. This is a hard failure with statically known domain bounds: GTIR type inference later hits assert expr_type.dims == target_type.dims in visit_SetAt (and under python -O, the mismatched IR reaches the backend).

Design

To decide whether a branch is ever selected, the pass makes the implicit broadcasts explicit (broadcast(a, (Vertex, K))) and reruns domain inference on the resulting concat_where. Since an explicit broadcast spans all dimensions of the expression, nothing is dropped from its inferred domain, and the pruning decision reduces to a plain branch.annex.domain.empty() check on that IR expression. Domain inference is used instead of computing the branch domains in the pass itself for simplicity and to not duplicate the domain semantics of concat_where. A surviving broadcast branch stays in the IR and is lowered by the subsequent RemoveBroadcast pass via the domain annex populated from the pruned concat_where — so both bugs disappear with one mechanism. Branches that are equal (up to the implicit broadcast) are pruned directly; the surviving branch is reinferred on the full domain of the concat_where, since its domains stem from the region where only its branch instance was selected.

To make the re-inference cheap, infer_domain.infer_expr gains a revisit_already_inferred option (default True, preserving behavior): with False, subexpressions that already carry a domain are neither descended into nor altered, so only the newly created broadcast nodes are populated.

Also in this PR:

• new ir_maker im.broadcast(expr, dims)
• infer_expr typed with a TypeVar so input and output expression types match

New test cases

Notation: concat_where(<cond>, <tb>: <tb dims>, <fb>: <fb dims>): <accessed domain> → <result>

Behavior changes (fail or crash against main's pass):

concat_where(K < 0,   a: [Vertex, K], b: [K]):      {Vertex: [0, 10), K: [0, 10)} → broadcast(b, (Vertex, K))
concat_where(K < 0,   a: [Vertex],    b: [Vertex]): {Vertex: [0, 10), K: [0, 10)} → broadcast(b, (Vertex, K))
concat_where(2 <= K,  a: [Vertex],    a: [Vertex]): {Vertex: [0, 10), K: [0, 10)} → broadcast(a, (Vertex, K))
concat_where(2 <= K,  a: [K],         broadcast(a, (Vertex, K))): {Vertex: [0, 10), K: [0, 10)} → broadcast(a, (Vertex, K))

Guard against the inverse mistake (a branch lacking the concat dimension is still selected via the implicit broadcast and must not be treated as never-selected):

concat_where(K < 5,   a: [Vertex],    b: [Vertex, K]): {Vertex: [0, 10), K: [0, 10)} → unchanged

Added coverage (already worked on main):

concat_where(K < 0,   a: [Vertex, K], b: [Vertex, K]): {Vertex: [0, 10), K: [0, 10)} → b
concat_where(10 <= K, a: [Vertex, K], b: [Vertex, K]): {Vertex: [0, 10), K: [0, 10)} → b
concat_where(K < 10,  a: [Vertex, K], b: [Vertex, K]): {Vertex: [0, 10), K: [0, 10)} → a
concat_where(0 <= K,  a: [Vertex, K], b: [Vertex, K]): {Vertex: [0, 10), K: [0, 10)} → a

The unit test additionally asserts, for every case, that the result's domain annex equals the accessed domain (RemoveBroadcast relies on it), and, for pruned broadcast results, that the domains inside the branch are inferred from the full domain of the concat_where.

Integration: a test for the field branch with fewer dimensions than the expression was added to test_concat_where.py:

concat_where(K < 0,   a: [I, J, K],   b: [K]):         full output domain → broadcast(b, (I, J, K))

Since statically known domain bounds are required for the pruning to trigger — before this PR no test compiled with static bounds, so the pass had no end-to-end coverage at all — the test is parametrized on a new static_domains fixture: dynamic_domains runs with runtime bounds as before, while static_domains creates the field operator with @gtx.field_operator(static_domains=True), making all field domains known at compile time so the pruning fires. The parametrization is extended to all tests in the file as it also improves coverage in general — in their static variant, the existing tests with never-selected branches now exercise the pass as well.

AI disclaimer: This PR was written with the help of AI tools. The description and pass implementation were reviewed in detail. Tests were only briefly reviewed, by putting them into a simple unified format and working with them makes me confident for them to be meaningful.

[havogt]

One case from #2767 is not covered by this PR, and since #2767 is now closed I want to flag it before it gets lost.

_infer_concat_where decides which region a branch is inferred on with

symbolic_cond if arg == true_field else cond_complement

infer_domain.py:382. That is value equality on IR nodes, so when the two branches are structurally equal, the second iteration also compares equal to true_field and the false branch is inferred on the condition instead of its complement. Both branches end up with the true-branch domain.

The result is a silent wrong answer, not a crash. Minimal reproducer, plain Cartesian, no connectivity:

@gtx.field_operator
def testee(a: IKField) -> IKField:
    t = a + 1.0
    return concat_where(KDim < 2, t, t)   # both branches equal, so the result is `t` everywhere

@gtx.program
def prog(a: IKField, out: IKField):
    testee(a, out=out, domain={IDim: (0, 3), KDim: (0, 4)})

with a all ones and out prefilled with -1, on run_gtfn:

[[ 2.  2. -1. -1.]
 [ 2.  2. -1. -1.]
 [ 2.  2. -1. -1.]]

K >= 2 is never written. The as_fieldop is emitted with domain KDim: [0, 2) while the SetAt target is KDim: [0, 4). Statically known bounds are required, as for the other two bugs.

	result
main 0d8fb34d3	wrong
this PR 808ed8f73	wrong
#2767	correct

So this is pre-existing on main rather than introduced here, but #2767 fixed it as part of the same change and this PR does not. The unit tests cannot catch it because they compare IR rather than execute; the closest case, concat_where(2 <= K, a, broadcast(a, (Vertex, K))), checks the structure but not the resulting domain.

The fix is to select by position instead of by value:

for is_true_branch, arg in [(True, true_field), (False, false_field)]:
    ...
    symbolic_cond if is_true_branch else cond_complement,

Applied to 808ed8f73 this makes the reproducer correct, and tests/next_tests/unit_tests/iterator_tests plus test_concat_where.py stay green, 737 passed on the non-GPU backends.

Otherwise the PR does what #2767 did, checked rather than read: #2767's regression test passes here on every backend and fails on main on all four compiled CPU backends, and both end-to-end examples behave. Handling the lower-dimensional branch by materialising the broadcast is also strictly more than #2767 did, which only declined to prune in that case.

[havogt]

lgtm, 2 structural comments. feel free to disagree. Might post an agent review as well in a bit.

[havogt]

Two issues in the new prune_empty_concat_where logic, both reproduced locally. The rest of the diff checks out - in particular the _infer_concat_where fix is a genuine latent bug (arg == true_field was true on both iterations when the branches were structurally equal, so identical branches both got the true-branch domain).

[havogt]

lgtm

[havogt]

Unrelated to this PR, but noticed while looking at concat_where pruning: concat_where(dim != i, a, b) silently evaluates to b everywhere.

InferDomainOps lowers != to an intersection where the semantics is a union — infer_domain_ops.py:91-99:

elif cpm.is_call_to(node, "not_eq"):
    # `IDim != a` -> `IDim < a & IDim > a`
    return self.visit(
        im.call("and_")(

{i : i != a} is {i < a} ∪ {i > a}, but {i < a} ∩ {i > a} is empty. canonicalize_domain_argument then faithfully expands the and_, and the true branch can never be selected:

out ← concat_where(c⟨ Iₕ: [-∞, 2[ ⟩, concat_where(c⟨ Iₕ: [3, ∞[ ⟩, a, b), b)

With a = ones, b = zeros on I: [0, 6):

concat_where(I != 2, a, b)      -> [0. 0. 0. 0. 0. 0.]
concat_where((I<2)|(I>2), a, b) -> [1. 1. 0. 1. 1. 1.]   # expected
concat_where(I == 2, a, b)      -> [0. 0. 1. 0. 0. 0.]   # control, correct

Only the compiled path is affected; Dimension.__ne__ with an integer raises NotImplementedError pointing at ADR 22.

This reproduces identically on main and on this branch — no regression from this PR, and the pruning changes here do not affect it either way. Filing it here only because it is in the same neighbourhood.

The fix looks like one token, im.call("and_") -> im.call("or_") at infer_domain_ops.py:94 together with the comment above it; note that test_infer_domain_ops.py:56-61 currently asserts the and_ shape, so it pins the current behaviour.