[CUDA][Reduce] Restore exact thread image counting - #2825
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
📝 WalkthroughWalkthroughThe PR updates the TVM submodule reference and changes scalar AllReduce thread-range resolution to use Z3 for exact, contiguous thread participation checks, with a bounding-span fallback when the prover returns unknown. ChangesAllReduce thread resolution
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ResolveAllReduceThreadRange
participant ReduceLayout
participant Z3
ResolveAllReduceThreadRange->>ReduceLayout: read thread expression and bounds
ResolveAllReduceThreadRange->>Z3: constrain and enumerate thread expression
Z3-->>ResolveAllReduceThreadRange: distinct thread count or unknown
ResolveAllReduceThreadRange->>ResolveAllReduceThreadRange: validate contiguous image or use bounding span
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![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 `@src/backend/common/op/reduce.h`:
- Around line 148-160: The scalar AllReduce lowering must not treat an unknown
distinct-thread count as proof that the bounding range is contiguous: update the
count handling around the existing count validation to fail lowering when the
result is unknown, unless an independent proof establishes every thread in
[base, end] participates. At src/backend/common/op/reduce.h lines 148-160,
remove the fallback that assigns span for unknown results; at lines 82-88,
update the documentation only if the implementation is intentionally changed to
accept an unproven bounding span.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: ae4ffb4c-b7ff-4f99-ac85-64590556d169
📒 Files selected for processing (2)
3rdparty/tvmsrc/backend/common/op/reduce.h
| if (count < 0) { | ||
| LOG(WARNING) << "tl.reduce: Z3 returned unknown while counting scalar " | ||
| "AllReduce participating threads; falling back to the " | ||
| "bounding range [" | ||
| << base << ", " << end << "]."; | ||
| count = span; | ||
| } | ||
| ICHECK_GT(count, 0) | ||
| << "tl.reduce: scalar AllReduce participating thread constraints are " | ||
| "unsatisfiable"; | ||
| ICHECK_EQ(count, span) | ||
| << "tl.reduce: partial scalar AllReduce requires one contiguous thread " | ||
| "range, but got " |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Reject an unknown result unless contiguity is proven.
Line 153 replaces the unknown distinct-thread count with span. Line 158 then always succeeds on this path. A sparse image, such as {base, base + 2}, can therefore configure a named barrier for three threads while only two threads execute the guarded reduction. The kernel can deadlock.
src/backend/common/op/reduce.h#L148-L160: Do not usespanas proof of contiguity. Fail lowering on an unknown result, or add an independent proof that every thread in[base, end]executes the reduction.src/backend/common/op/reduce.h#L82-L88: Update the documentation if the implementation can accept an unproven bounding span.
📍 Affects 1 file
src/backend/common/op/reduce.h#L148-L160(this comment)src/backend/common/op/reduce.h#L82-L88
🤖 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/backend/common/op/reduce.h` around lines 148 - 160, The scalar AllReduce
lowering must not treat an unknown distinct-thread count as proof that the
bounding range is contiguous: update the count handling around the existing
count validation to fail lowering when the result is unknown, unless an
independent proof establishes every thread in [base, end] participates. At
src/backend/common/op/reduce.h lines 148-160, remove the fallback that assigns
span for unknown results; at lines 82-88, update the documentation only if the
implementation is intentionally changed to accept an unproven bounding span.
1 participant
First see & merge tile-ai/tvm#61 .
Created safe fallback for counting threads in
src/backend/common/op/reduce.hwhen determinating named barrier count.Summary
unknown.3rdparty/tvmsubmodule. Mergetile-ai/tvm#61first.C++ style / lint notes
src/backend/common/op/reduce.h.docs/developer_guide/cpp_style.md.