fix[next]: decide concat_where pruning from the concat_where's domain - #2767
fix[next]: decide concat_where pruning from the concat_where's domain#2767havogt wants to merge 2 commits into
Conversation
`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`.
There was a problem hiding this comment.
Pull request overview
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
This PR fixes prune_empty_concat_where to decide branch pruning using the concat_where expression’s own (inferred) domain, preventing incorrect “too weak/too strong” pruning and centralizing branch-domain computation in domain_utils.
Changes:
- Introduces
domain_utils.concat_where_branch_domain()and reuses it from both domain inference and the prune transform. - Updates
prune_empty_concat_whereto only prune when the surviving branch covers allconcat_wheredimensions (avoids silent dimension dropping). - Adds regression tests covering missing concat dimension, dimension dropping, and complement-before-promote ordering.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_prune_empty_concat_where.py |
Adds regression tests for correct pruning behavior across dimension-mismatch scenarios. |
tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py |
Adds a unit test ensuring complement is taken before promote for concat_where branch domains. |
src/gt4py/next/iterator/transforms/prune_empty_concat_where.py |
Reworks prune logic to use concat_where domain-derived branch regions and guard against dimension dropping. |
src/gt4py/next/iterator/transforms/infer_domain.py |
Refactors _infer_concat_where to delegate branch-region computation to domain_utils. |
src/gt4py/next/iterator/ir_utils/domain_utils.py |
Adds concat_where_branch_domain() helper encapsulating complement-before-promote ordering. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| 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() | ||
| ) |
| 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 |
| 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. |
| # the other order does not work | ||
| with pytest.raises(AssertionError): | ||
| domain_utils.domain_complement(domain_utils.promote_domain(cond, domain.ranges.keys())) |
|
I haven't read the code and essentially stopped at
I think it is not needed to couple the is the same as I'll add further comments after reading the code. |
Thanks — I implemented your suggestion to see where it leads, and it uncovered something worth reporting before the design is settled: taken literally, pruning from the branch domains cannot fix the first of the two bugs, because the information is not there. The counterexampleTwo expressions, same condition, same branch types, and byte-identical branch domains, that require opposite answers:
with The reason is What that leavesThree ways to complete the design, as far as I can see:
Option 2 is what your own comment writes out ( One thing that argues against the framingThe central promotion helper you half-remembered does exist: So the one pass that already promotes Severity, since it affects how much this is worthThe two bugs are not equally serious.
Cost, which is the part I am least sure about
The two commits are split so the second can be dropped, in which case the never-selected-branch bug stays open and cannot be closed without either The branch is at havogt#76 if you want to look at it; happy to fold whichever shape you prefer into #2767. |
|
Do you have an end-to-end example that fails without this fix? I'm currious what happens in that case. |
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.
| ) | ||
|
|
||
|
|
||
| def test_concat_where_never_selected_branch_keeps_dimensions(cartesian_case): |
There was a problem hiding this comment.
wondering if we should make the cases to parametrize over static and runtime domains (or at least make it selectable). I only realize now that probably we only have very limited tests with static domains, but that's our production use-case.
|
Asked it to add a frontend test for the failing case. Note that I initially worked on a fix for the second case, when the prune is not done but could be done. Do you want an example for that as well? |
prune_empty_concat_wheredecided whether a branch is ever selected from the branch's own domain, filtered to the dimensions of the condition. That is the wrong quantity, and it is wrong in both directions.Too weak
A branch's 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 an empty mapping is not an empty domain. A branch selected nowhere survives.
over
Vertex: [0, 10),K: [0, 10).ais selected whereK < 0, i.e. nowhere, so the expression isb— buta's domain isVertex: [0, 10)with noKrange to be empty, so it survived.Too strong
Nothing checked that the surviving branch spans the same dimensions as the
concat_where, so witha: Field[[Vertex, K]]andb: Field[[K]]the same expression is rewritten tob, replacing a two dimensional expression by a one dimensional one. This one reproduces onmaintoday, independent of the case above.The fix
Both follow from deriving the selected region from the
concat_where's own domain — which is exactly whatinfer_domain._infer_concat_wherealready computes per branch. That computation moves todomain_utils.concat_where_branch_domainand both callers use it, so the requirement to take the complement before promoting —domain_complementis undefined on a range infinite on both sides — exists in one place rather than being duplicated.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 the pipeline does not set:_infer_concat_whererebuilds the node throughim.call, which produces an untypedFunCall, sonode.typeisNoneat everyconcat_wherethis pass sees.The pass body goes from 90 lines to 46; the three overlapping mechanisms it had — an
empty()early return, the branch domain filter, and a type comparison — collapse to one.Testing
Three regression tests, each verified to fail against the previous implementation: the never selected branch that lacks the concat dimension, the dimension dropping prune, and a
NEVERdomain that previously crashed. Plus the complement before promote ordering, tested intest_domain_utils.pywhere the helper now lives.504
iterator_testspass, 39 doctests, and 304 CPU integration tests acrosstest_concat_where.py/test_where.py/test_named_collections.py.ruffandmypyclean.Not behaviour preserving, deliberately
Two intentional differences: pruning is declined where the old code would drop a dimension, and input with no
annex.domaintogether with an empty condition now raises rather than pruning — the pass documents that it requires inferred domains, and the doctest was the only such caller. The doctest now runsinfer_exprfirst.