tvm
tvm copied to clipboard
[Fix][TIR] LowerThreadAllreduce warp reduction mask
The warp reduction implemented by "shuffle down" primitive takes a mask denoting the active threads within the warp that participate in this shuffle.
Previously we compute the mask, while in practice we find that it results in "CUDA illegal instruction" error on NVIDIA H100 GPU when the mask is set, and the issue is gone if we do not update the mask. Therefore, this PR updates the allreduce lowering to remove the mask update.
Confirmed the correctness on the following devices:
- NVIDIA H100,
- NVIDIA RTX 4090,
- AMD Radeon 7900 XTX,
- Apple M2 Ultra.