[Bugfix] Make ThreadAllReduce pass compatible with int64

[yzh119]

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:

  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.

cc @vinx13 @junrushao @Hzfengsy

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• No users to tag found in teams: bugfix _{See #10317 for details}

_{Generated by tvm-bot}