fix(ir): make composite collective barrier signals reusable - #2175
fix(ir): make composite collective barrier signals reusable#2175georgebisbas wants to merge 2 commits into
Conversation
The composite collective lowerings shared a `Set(1)` + `Wait(>= 1)` completion barrier. That leaves every signal cell at `1` after a call returns, so a later invocation reusing the same signal could satisfy its waits from stale state before peers finished the new data transfers, returning mixed or partial results. Replace it with one generation protocol shared by the whole `pld.tensor.*` family: each barrier does `AtomicAdd(1)` into every peer's cell, then `Wait(>= generation)`, where `generation` counts how many barriers have been issued against that signal. Cells become a monotonic per-rank barrier count, so the n-th barrier waits for `>= n` and a signal is safely reusable across back-to-back calls — including chains that mix collectives. Generations are assigned at lowering time by `SignalGenerationTable`, keyed on the signal's allocation identity so the rebind idiom `sig = pld.tensor.barrier(sig)` keeps counting on one counter instead of restarting. Collectives barrier through `LoweringBuilder::EmitBarrier`, so a new collective inherits the protocol rather than rolling its own. This also removes a latent hazard in `reduce_scatter`, which mixed `Set` and `AtomicAdd` on the same cells — a set can clobber an already-advanced counter. Cases the compile-time scheme cannot make safe are now rejected with an actionable error instead of silently degrading: - any collective inside a dynamic `for` / `while` loop (previously only `allreduce` was rejected; the other five were silently broken) - one signal shared between the per-cell mesh and per-row ring protocols - reuse after a mesh `allreduce` over a symbolic extent, whose per-chunk barrier count is not statically known Tests: protocol assertions (generation sequence, AtomicAdd/Ge, rebind chaining, rejections) in `test_lower_composite_ops.py`, plus a back-to-back distributed ST for `all_to_all` / `allgather` / `reduce_scatter` on a shared signal. Closes hw-native-sys#2156
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe composite lowering pass now assigns reusable compile-time barrier generations across tensor collectives. Lowerings, DSL and developer documentation, unit tests, and distributed tests were updated for AtomicAdd/Ge synchronization, signal aliasing, reuse restrictions, and loop validation. ChangesCollective barrier signal reuse
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant CollectiveLowering
participant SignalGenerationTable
participant DistributedSignal
CollectiveLowering->>SignalGenerationTable: allocate next generation
CollectiveLowering->>DistributedSignal: AtomicAdd notification
DistributedSignal-->>CollectiveLowering: Ge wait reaches generation
CollectiveLowering->>SignalGenerationTable: reserve later generations
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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
Here are some automated review suggestions for this pull request.
Reviewed commit: e6219ce9ea
ℹ️ 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".
| if (!IsTensorCollective(call)) return; | ||
| CHECK_SPAN(repeating_scope_depth_ == 0, call->span_) | ||
| << "pld.tensor.allreduce is not supported inside a for/while loop. " | ||
| "The signal protocol is single-use and cannot reuse a signal across dynamic invocations."; | ||
| << call->op_->name_ |
There was a problem hiding this comment.
When a pld.tensor.* collective appears inside an if/else, the mutator still lowers both branches while using a single compile-time generation counter, but only one branch executes at runtime. For example, if flag: sig = pld.tensor.barrier(sig) else: sig = pld.tensor.barrier(sig) assigns generation 1 to the then branch and generation 2 to the else branch, so taking the else branch deadlocks waiting for a counter that was incremented only once; a later barrier after the if can also wait on stale state because the IfStmt result is not tied to the executed branch's generation. The new check only forbids for/while, so conditional collectives remain silently miscompiled; either reject them or model generation as control-flow-dependent.
AGENTS.md reference: AGENTS.md:L36-L39
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in the latest push. The repeating_scope_depth_ counter has been renamed to scope_depth_ and now also increments for IfStmtPtr visitors, so collectives inside for/while/if/else are all rejected. The error message has been updated accordingly: "is not supported inside a for/while/if/else scope."
Additionally, LoweringBuilder now tracks a nested_ flag and guards EmitBarrier() with INTERNAL_CHECK_SPAN(!nested_) — inner builders created by EmitFor, EmitIf, and EmitIfExpr are all marked nested so they cannot emit barriers directly.
![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
🧹 Nitpick comments (1)
src/ir/transforms/lower_composite_ops_pass.cpp (1)
444-457: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider enforcing the "straight-line only" contract for
EmitBarrier.The restriction is documented but unchecked; calling it from an
EmitFor/EmitIfbody would reserve one generation for a barrier that executes N times, producing a hang or a silently degraded barrier. Abool nested_flag set on builders constructed by the control-flow emitters plus anINTERNAL_CHECK_SPAN(!nested_, span)here would make the misuse a build-time failure rather than a runtime deadlock.🤖 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 `@src/ir/transforms/lower_composite_ops_pass.cpp` around lines 444 - 457, Enforce EmitBarrier’s straight-line-only contract by tracking whether the builder is nested in control flow: set a nested_ flag on builders created by EmitFor/EmitIf, then add INTERNAL_CHECK_SPAN(!nested_, span) at the start of EmitBarrier. Ensure top-level builders remain allowed while nested calls fail during construction.
🤖 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 `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 1831-1845: Extend CheckCollectiveLoopUse to reject tensor
collectives when either loop or conditional scope depth is nonzero. Track IfStmt
branch entry and exit with a conditional-scope depth counter in the mutator, and
include that counter in the existing CHECK_SPAN condition while preserving the
current loop diagnostic behavior.
---
Nitpick comments:
In `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 444-457: Enforce EmitBarrier’s straight-line-only contract by
tracking whether the builder is nested in control flow: set a nested_ flag on
builders created by EmitFor/EmitIf, then add INTERNAL_CHECK_SPAN(!nested_, span)
at the start of EmitBarrier. Ensure top-level builders remain allowed while
nested calls fail during construction.
🪄 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: 84dbabc7-7cca-4f18-bb23-fdf980647950
📒 Files selected for processing (10)
docs/en/dev/distributed_ops.mddocs/en/dev/passes/12-lower_composite_ops.mddocs/en/dev/passes/38-synthesize_allreduce_signals.mddocs/zh-cn/dev/distributed_ops.mddocs/zh-cn/dev/passes/12-lower_composite_ops.mddocs/zh-cn/dev/passes/38-synthesize_allreduce_signals.mdpython/pypto/language/distributed/op/tensor_ops.pysrc/ir/transforms/lower_composite_ops_pass.cpptests/st/distributed/collectives/test_l3_tensor_collective_signal_reuse.pytests/ut/ir/transforms/test_lower_composite_ops.py
e6219ce to
e85c587
Compare
georgebisbas
left a comment
georgebisbas
left a comment
There was a problem hiding this comment.
All AI review feedback addressed in the latest push:
-
Conditional scope tracking: Renamed
repeating_scope_depth_toscope_depth_, addedVisitStmt_(const IfStmtPtr&)override.CheckCollectiveLoopUsenow rejects collectives insidefor/while/if/else. -
Nested barrier guard: Added
nested_flag toLoweringBuilder, nested constructor, andINTERNAL_CHECK_SPAN(!nested_)at the top ofEmitBarrier(). All inner builders fromEmitFor/EmitIf/EmitIfExprare now marked nested. -
Clang-tidy fix: Removed unused
#include "pypto/ir/program.h"
…ys#2156 Move pld.system.notify/wait docs before Signal lifetime for readability. Add test_unrolled_loop_collective_gets_fresh_generations to verify the pl.unroll escape hatch for the loop rejection added in the protocol. Route operator names through ir.get_op() so typos raise at import.
e85c587 to
289046e
Compare
1 participant
Closes #2156.
The composite collective lowerings shared a
Set(1)+Wait(>= 1)completion barrier. That leaves every signal cell at1after a call returns, so a later invocation reusing the same signal could satisfy its waits from stale state before peers finished the new data transfers — returning mixed or partial results. This replaces it with one generation protocol shared by the wholepld.tensor.*family, so signals are safely reusable.The protocol
AtomicAddturns each cell into a monotonically increasing count of how many barriers the rank owning that row has entered, so the n-th barrier issued against a signal waits for>= n.Latent bug fixed on the way
reduce_scattermixedSet(ready barrier) andAtomicAdd(post-reduce barrier) on the same cells. Every collective now notifies withAtomicAddonly.New compile-time rejections (previously silent corruption)
for/whileloop (all six, not justallreduce)allreduceover a symbolic extentChanges
src/ir/transforms/lower_composite_ops_pass.cpp—SignalGenerationTable+LoweringBuilder::EmitBarrierhelpers; all six collectives convertedtests/ut/ir/transforms/test_lower_composite_ops.py— generation sequencing,AtomicAdd/Ge, rebind chaining,pl.unrollescape hatch, rejections (80 passed)docs/en/dev/distributed_ops.md+docs/zh-cn/dev/distributed_ops.md— notify/wait docs restructuredTest plan
test_lower_composite_ops.py, 2622 wider UT regressiontest_l3_tensor_collective_signal_reuse.py— back-to-backall_to_all,allgather,reduce_scattergreen on sim (P=2)