Skip to content

Commit

Permalink
[TIR] Output DeclBuffer in LowerThreadAllreduce (#15646)
Browse files Browse the repository at this point in the history
For any local buffers allocated as part of `LowerThreadAllreduce`,
generate a `DeclBuffer`.

This is a subset of the changes made in
#14778, broken out for ease of
testing and review.
  • Loading branch information
Lunderberg authored Sep 11, 2023
1 parent 25b8a07 commit 4d7e93c
Show file tree
Hide file tree
Showing 2 changed files with 201 additions and 231 deletions.
1 change: 1 addition & 0 deletions src/tir/transforms/lower_thread_allreduce.cc
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
// Fix all local allocations as all statements are built.
Stmt body = SeqStmt::Flatten(seq);
for (Buffer buf : new_alloc_bufs) {
body = DeclBuffer(buf, body);
body = Allocate(buf->data, buf->dtype, buf->shape, const_true(buf->dtype.lanes()), body);
}

Expand Down
Loading

0 comments on commit 4d7e93c

Please sign in to comment.