feat(ir): dispatch-predicate scope form for pl.at - #2151
Conversation
Refs hw-native-sys#2066 Extends `predicate=` from `pl.submit` / `pl.spmd_submit` (Phase 1, hw-native-sys#2067) and the `with pl.spmd(...)` scope form (Phase 3, hw-native-sys#2092) to `pl.at`: with pl.at(level=pl.Level.CORE_GROUP) as g_tid: rc = pl.store(pl.load(rc, [0, 0], [128, 128]), [0, 0], rc) with pl.at(level=pl.Level.CORE_GROUP, deps=[g_tid], predicate=(rc[0, 0] > 0)) as tid: out = pl.store(pl.load(x, [0, 0], [128, 128]), [0, 0], out) Same expression, validation, lowering and runtime ABI as the existing forms — the predicate rides on `InCoreScopeStmt::attrs_` until `OutlineIncoreScopes` moves it onto `Submit::predicate_`. The scope-attr walkers (visitor, mutator, ConvertToSSA, DCE, structural_hash, NoNestedCall) were already scope-kind agnostic, so this is parser + dsl_api + printer, as the tracking issue predicted. Placement is restricted to where the scope actually becomes an independently submitted L0 task, since anywhere else the predicate has no runtime carrier and would be silently dropped: * `level=pl.Level.CORE_GROUP` only — a Hierarchy `pl.at` inside an Orchestration function survives the whole Default pipeline un-outlined (`OutlineHierarchyScopes` is Opaque-only) and never produces a `Submit`. * not nested inside `pl.cluster()` / `pl.spmd()` / another `pl.at` — the inner dispatch is folded into the enclosing Group / Spmd wrapper, and codegen emits `set_predicate` from the outer call only. Both are rejected at parse time, with `ScopeOutliner` and a new `AssertNoInnerDispatchPredicate` in `OutlineClusterScopes` re-asserting them for hand-built / deserialized IR (mirroring the `UnwrapNestedSpmd` guard). Scope-producer tracking (`_record_scope_producer`, added for `pl.spmd` in Phase 3) now also runs on `with pl.at(...) as tid:`, so the producer-in-`deps=` contract genuinely applies. Unlike `pl.spmd`, `pl.at` accepts `deps=` on every form, so the remediation hint always offers it. Outliner fix: `ScopeOutliner` read the predicate attr raw, bypassing its own `store_target_renames_` substitution. A `pl.at` body typically writes its tensor with `pl.store`, so the attr names a scope-local post-store alias that outlining drops — the emitted `Submit` then referenced a dangling Var (UseAfterDefCheck). The operand is now resolved exactly as the synthesised call's args are. This also closes the same latent hole for a `pl.spmd` predicate over a tensor written by an earlier sibling `pl.at` scope. Tests: parser (attr landing, canonical order, print round-trip, no-`as tid` form, level + three nesting rejections, shared shape validation, scope-producer contract, hint wording), transforms (SSA versioning, outliner rebinding, DCE liveness), codegen (`set_predicate` block, operator/operand-order, emission order, and a `windowize=True` case — the only DSL path that reaches the predicate entry in WindowExternalization's view-synthesised key filter). On-board a2a3 ST gains a `pl.at` variant of the predicated-dispatch chain: 6/6 pass across all three authoring forms. Full UT suite 7791 passed, 2 skipped. Docs: `docs/en` + `docs/zh-cn` python_syntax — the stale "not accepted on pl.at" note is replaced with the surface, its two placement rules, and the `deps=` difference vs `pl.spmd`.
📝 WalkthroughWalkthroughAdds ChangesCORE_GROUP
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant User as Python DSL
participant Parser as AST parser
participant Scope as InCoreScopeStmt
participant Outliner as ScopeOutliner
participant Submit as Submit
participant Codegen as Runtime codegen
User->>Parser: define pl.at(predicate=...)
Parser->>Scope: validate and attach predicate
Scope->>Outliner: provide scope predicate
Outliner->>Submit: rewrite operands and transfer predicate
Submit->>Codegen: emit set_predicate before submit
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Check pre-existing wrapper functions for inner predicates
When the input program already contains a FunctionType::Group or FunctionType::Spmd function with a predicated inner Submit, the main pass loop places that function in new_functions and continues, while this guard iterates only all_outlined_functions. Such parser-authored, hand-built, or deserialized wrappers therefore bypass AssertNoInnerDispatchPredicate, and the outer wrapper dispatch can silently ignore the inner predicate; run the guard over both existing and newly outlined wrapper functions.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for kind, owner in ( | ||
| (ir.ScopeKind.Cluster, "pl.cluster()"), | ||
| (ir.ScopeKind.Spmd, "pl.spmd()"), | ||
| (ir.ScopeKind.InCore, "another pl.at(level=pl.Level.CORE_GROUP)"), | ||
| ): |
There was a problem hiding this comment.
When a predicated CORE_GROUP pl.at is nested in a non-CORE_GROUP pl.at, the scope stack contains ScopeKind.Hierarchy, which this list never checks; similarly, a top-level predicated pl.at inside an explicitly declared Group or Spmd function is accepted even though OutlineIncoreScopes skips those function types. In both cases the scope does not become an independently dispatched L0 task and its predicate can be dropped, despite the updated documentation prohibiting nesting inside another pl.at; include the hierarchy and enclosing-function contexts in the placement gate.
AGENTS.md reference: AGENTS.md:L38-L39
Useful? React with 👍 / 👎.
| // ``pl.at(level != CORE_GROUP, predicate=...)``; guard here for hand-built / | ||
| // deserialized IR, mirroring the UnwrapNestedSpmd guard in | ||
| // OutlineClusterScopes. | ||
| INTERNAL_CHECK_SPAN(!scope_predicate || !As<HierarchyScopeStmt>(op), op->span_) |
There was a problem hiding this comment.
For hand-built or deserialized IR containing a predicated HierarchyScopeStmt in an Orchestration function, OutlineHierarchyScopes skips the function because it only processes Opaque functions, so this OutlineScope check is never reached. The unsupported scope therefore survives the pipeline with its predicate instead of triggering the promised backstop, allowing codegen to silently lose the dispatch condition; perform this validation before the function-type skip or in a program-wide verifier.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/ut/ir/transforms/test_submit_passes.py`:
- Around line 530-534: Strengthen the assertions after printing `main` to
inspect the outlined `Submit` nodes and verify that the predicated submit’s
operand is the call-site `rc` result exported by the preceding gate submit. Keep
the existing predicate and `__FREE_VAR` checks, but distinguish the gate
submit’s exported result from the pre-store/input `rc`.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 8f060f2f-c6de-495f-9586-b2c0f4abede6
📒 Files selected for processing (11)
docs/en/dev/language/00-python_syntax.mddocs/zh-cn/dev/language/00-python_syntax.mdinclude/pypto/ir/transforms/utils/scope_outline_utils.hpython/pypto/language/dsl_api.pypython/pypto/language/parser/ast_parser.pysrc/ir/transforms/outline_cluster_scopes_pass.cppsrc/ir/transforms/python_printer.cpptests/st/runtime/scheduling/test_predicated_dispatch.pytests/ut/codegen/test_predicate_codegen.pytests/ut/ir/transforms/test_submit_passes.pytests/ut/language/parser/test_submit_predicate.py
| main = program.get_function("main") | ||
| assert main is not None | ||
| printed = ir.python_print(main) | ||
| assert "predicate=(" in printed, printed | ||
| assert "__FREE_VAR" not in printed, f"predicate references a dangling Var:\n{printed}" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert the predicate binds to the gate submit’s exported rc result.
These checks also pass if outlining binds the predicate to the pre-store/input rc: it is defined and prints without __FREE_VAR, but dispatch observes the wrong value. Inspect the outlined Submit nodes and assert the predicated submit’s operand is the call-site rc result exported by the preceding gate submit.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/ut/ir/transforms/test_submit_passes.py` around lines 530 - 534,
Strengthen the assertions after printing `main` to inspect the outlined `Submit`
nodes and verify that the predicated submit’s operand is the call-site `rc`
result exported by the preceding gate submit. Keep the existing predicate and
`__FREE_VAR` checks, but distinguish the gate submit’s exported result from the
pre-store/input `rc`.
1 participant
What this does
Extends
predicate=towith pl.at(level=pl.Level.CORE_GROUP, ...):Same expression, validation, lowering and runtime ABI as the two existing forms. The predicate rides on
InCoreScopeStmt::attrs_from parse untilOutlineIncoreScopesmoves it ontoSubmit::predicate_.Why it really was "parser + dsl_api only"
Phase 3's prediction held. Every scope-attr walker —
IRVisitor::VisitScopeAttrs,IRMutator::MutateScopeAttrs,ConvertToSSA::SubstScopeAttrs, DCEFindLiveRoots,structural_hash,NoNestedCall::ShouldVisitScopeAttr— andScopeOutlineritself are already scope-kind agnostic.pl.atInCore scopes are outlined at pass 8, a strict subset of the passes thepl.spmdpredicate (pass 9) already survives, so there is no new pass exposure. The only C++ additions are the printer hookup and two backstop asserts.Placement restrictions
A predicate only has a runtime carrier where the scope becomes an independently submitted L0 task. Two placements would silently drop it, so both are rejected at parse time:
level != CORE_GROUPHierarchyScopeStmt.OutlineHierarchyScopesonly rewrites scopes insideOpaquefunctions, so inside an@pl.function(type=Orchestration)the scope survives the entireDefaultpipeline un-outlined and never produces aSubmit(verified empirically).pl.cluster()/pl.spmd()/ anotherpl.atBuildSpmdCallDispatchPlan/BuildAivOnlyGroupDispatchPlan/BuildMixedGroupDispatchPlanall pass the outer wrapper call toEmitPredicateHint.Mirrors
_reject_spmd_submit_only_kwargs_in_cluster. For hand-built / deserialized IR both are re-asserted in C++ —ScopeOutlinerguards the Hierarchy case, and a newAssertNoInnerDispatchPredicateinOutlineClusterScopeswalks each outlined Group / Spmd body for aSubmit::predicate_(same pattern as the existingUnwrapNestedSpmdguard).Scope-producer tracking
_record_scope_producer(added forpl.spmdin Phase 3) now also runs onwith pl.at(...) as tid:, so the Phase-2 producer∈deps=contract genuinely applies to tensors apl.atbody writes. Unlikepl.spmd,pl.atacceptsdeps=on every form, so the remediation hint always offersdeps=[tid]instead of redirecting to theas tidcapture form (which would be wrong advice here).Outliner bug found and fixed (also affects the merged
pl.spmdform)ScopeOutlinerreadkAttrPredicateraw, bypassing its ownstore_target_renames_substitution — a target-kind scope skips the baseMutateScopeAttrswalk, sinceVisitScopeKindgoes straight toOutlineScope.A
pl.atbody typically writes its tensor withpl.store, so after SSA the attr names a scope-local post-store alias. Outlining exports that store target under a fresh call-site Var and drops the alias, leaving the emittedSubmitreferencing a Var with no definition in the parent function:UseAfterDefCheckcatches it, so it fails loudly rather than miscompiling — but it made the feature unusable for the most naturalpl.atspelling. The operand is now resolved through the same rename map the synthesised call's args use, i.e. to the value current as of this scope.This was latent for
pl.spmdtoo: apl.spmdpredicate over a tensor written by an earlier siblingpl.atscope hits the identical path. Thepl.spmdSTs did not surface it because their producer scopes bind via a call result, not a store.Testing
test_submit_predicate.py, +12 tests) — attr landing, canonical attr order, print → reparseassert_structural_equal, no-as tidform, the level rejection, all three nesting rejections, shared shape validation, the scope-producer contract (accepts withdeps=, rejects without), and the hint wording.test_submit_passes.py) — SSA versions the operand inside the InCore scope attr;test_outlining_rebinds_a_store_produced_predicate_operandpins the fix above (it reproduces the__FREE_VARoutput without it); DCE keeps a Var used only in the predicate index.test_predicate_codegen.py, +9 tests) — theL0TaskPredicateblock with&ext_rc, operator/operand-order mapping,set_predicatebeforert_submit_*, exactly one predicated task. Plus awindowize=Truecase:windowizeispl.at-only, making this the only DSL path that reaches thekAttrPredicateentry inWindowExternalization's view-synthesised key filter, which until now was written defensively and unreachable from the frontend. Verified the rewrite genuinely fires (the emitted args switch toext_x.view(...)).tests/st/runtime/scheduling/test_predicated_dispatch.pygains apl.atvariant with the identical task chain and expected values. 6/6 pass across all three authoring forms (gate=0 skips, gate=1 dispatches; the consumer unlocks in both cases, so inline-retire still settles fanout).Docs
docs/en/dev/language/00-python_syntax.md+docs/zh-cn/...— the stale "It is not accepted onpl.at(...)" note is replaced with the surface, apl.atsubsection covering the two placement rules, and thedeps=-on-every-form difference vspl.spmd. Thepl.at ... as tidtable row is updated to match howpl.spmd's row documents the kwarg.Refs #2066