This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c98e29bbf4 [Bugfix] Make ThreadAllReduce pass compatible with int64
(#14991)
c98e29bbf4 is described below
commit c98e29bbf481c023c4fc979eeb25b5484bac1e1e
Author: Zihao Ye <[email protected]>
AuthorDate: Wed May 31 09:31:03 2023 -0700
[Bugfix] Make ThreadAllReduce pass compatible with int64 (#14991)
# The Issue
Currently, the ThreadAllReduce pass would throw an error when the mask data
type is uint32 and the `group_index`'s data type is int64:
```bash
1: tvm::tir::ThreadAllreduceBuilder::MakeAllreduce(tvm::tir::CallNode
const*)
at
/home/zhye/repos/relax/src/tir/transforms/lower_thread_allreduce.cc:362
0: tvm::tir::BufferStore::BufferStore(tvm::tir::Buffer, tvm::PrimExpr,
tvm::runtime::Array<tvm::PrimExpr, void>, tvm::Span)
at /home/zhye/repos/relax/src/tir/ir/stmt.cc:477
File "/home/zhye/repos/relax/src/tir/ir/stmt.cc", line 477
TypeError: dtype mismatch on BufferStore: buffer's dtype is `uint32`, the
lanes of indexing are: `1`, but RHS's dtype is `int64`
```
As int64 becomes the standard index data type for large models, we should
fix the issue.
# The Fix
This PR resolves the issue by casting the `group_index` to the data type
used in mask.
---
src/tir/transforms/lower_thread_allreduce.cc | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/src/tir/transforms/lower_thread_allreduce.cc
b/src/tir/transforms/lower_thread_allreduce.cc
index 9104b7d51b..5c004fa5db 100644
--- a/src/tir/transforms/lower_thread_allreduce.cc
+++ b/src/tir/transforms/lower_thread_allreduce.cc
@@ -355,7 +355,8 @@ class ThreadAllreduceBuilder final : public StmtExprMutator
{
{
PrimExpr mask = Call(mask_dtype, builtin::tvm_warp_activemask(), {});
if (group_extent > 1) {
- mask = mask & (((1 << reduce_extent) - 1) << (reduce_extent *
group_index));
+ mask = mask &
+ (((1 << reduce_extent) - 1) << (reduce_extent *
cast(mask_dtype, group_index)));
}
seq.emplace_back(BufferStore(mask_buffer, mask, zero_indices));
// Push the buffer description. Later this will have an
