Repository navigation
[BugFix] Add FirstOwningBlockCollector to map allocated buffers to their declaring blocks - #2007
Rachmanino wants to merge 2 commits into
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! 🚀 |
📝 WalkthroughWalkthroughThis change introduces a new Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~35 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
…ring blocks This change introduces the FirstOwningBlockCollector class, which collects the first Block that declares each allocated buffer variable. The BufferAllocationLocator is updated to utilize this new collector, ensuring that shared barriers remain associated with their declaring blocks during buffer allocation. This enhancement improves the management of buffer allocations and addresses potential issues with opaque child blocks in pipelined For loops.
![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 the current code and only fix it if needed.
Inline comments:
In `@src/transform/plan_update_buffer_allocation_location.cc`:
- Around line 178-188: The constructor currently pushes barrier buffers into
alloc_buffers_ (using managed_allocations_, owning_block and alloc_buffers_)
which causes duplicates because VisitStmt_(const BlockNode*) already preserves
barrier buffers from op->alloc_buffers; change the constructor logic to skip any
buffer where buffer.scope() indicates a barrier (or where
IsBarrierBuffer(buffer) would be true) so you do not insert barrier buffers into
alloc_buffers_ (leave VisitStmt_ to handle them), i.e., add a guard before
alloc_buffers_[ob->second].push_back(buffer) to continue if the buffer is a
barrier or remove that push entirely for barrier-scoped buffers.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 47ff5ca2-7020-4d47-95ae-dcd48fe0307b
📒 Files selected for processing (1)
src/transform/plan_update_buffer_allocation_location.cc
| if (buffer.scope() == "shared.barrier" || | ||
| buffer.scope() == "shared.cluster_barrier") { | ||
| if (managed_allocations_.count(buffer->data.get())) { | ||
| auto ob = owning_block.find(buffer->data.get()); | ||
| ICHECK(ob != owning_block.end()) | ||
| << "PlanAndUpdateBufferAllocationLocation: barrier buffer " | ||
| << buffer->name << " has no owning Block"; | ||
| alloc_buffers_[ob->second].push_back(buffer); | ||
| } | ||
| continue; | ||
| } |
There was a problem hiding this comment.
Critical: Barrier buffers will be allocated twice in their owning block.
The constructor now adds barrier buffers to alloc_buffers_[owning_block] (line 185). However, VisitStmt_(const BlockNode*) at lines 294-300 also preserves barrier buffers from the original op->alloc_buffers:
for (const Buffer &buf : op->alloc_buffers) {
if (IsBarrierBuffer(buf)) {
alloc_buffers.push_back(buf); // First addition
...
}
}
auto it = alloc_buffers_.find(op);
if (it != alloc_buffers_.end()) {
for (const Buffer &buf : it->second) {
alloc_buffers.push_back(buf); // Second addition (same barrier buffer)
...
}
}When visiting the owning block, the same barrier buffer is added from both sources, resulting in duplicate allocations.
Proposed fix: Skip barrier buffers from alloc_buffers_ since they're already preserved from op->alloc_buffers
auto it = alloc_buffers_.find(op);
if (it != alloc_buffers_.end()) {
for (const Buffer &buf : it->second) {
+ // Barrier buffers are already preserved from op->alloc_buffers above
+ if (IsBarrierBuffer(buf)) {
+ continue;
+ }
alloc_buffers.push_back(buf);
PushBinding(buf->data, buf);
}
}Alternatively, revert to the original behavior where barrier buffers are not added to alloc_buffers_ at all, since VisitStmt_ already preserves them from their declaring block's op->alloc_buffers.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/transform/plan_update_buffer_allocation_location.cc` around lines 178 -
188, The constructor currently pushes barrier buffers into alloc_buffers_ (using
managed_allocations_, owning_block and alloc_buffers_) which causes duplicates
because VisitStmt_(const BlockNode*) already preserves barrier buffers from
op->alloc_buffers; change the constructor logic to skip any buffer where
buffer.scope() indicates a barrier (or where IsBarrierBuffer(buffer) would be
true) so you do not insert barrier buffers into alloc_buffers_ (leave VisitStmt_
to handle them), i.e., add a guard before
alloc_buffers_[ob->second].push_back(buffer) to continue if the buffer is a
barrier or remove that push entirely for barrier-scoped buffers.
|
would you mind providing a test to reproduce? |
|
closed as already fixed in pr #1971 |
This fixes missing barrier init attr when TMA is disabled.
This change introduces the FirstOwningBlockCollector class, which collects the first Block that declares each allocated buffer variable. The BufferAllocationLocator is updated to utilize this new collector, ensuring that shared barriers remain associated with their declaring blocks during buffer allocation. This enhancement improves the management of buffer allocations and addresses potential issues with opaque child blocks in pipelined For loops.
Summary by CodeRabbit