tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[Fix][TIR] LowerThreadAllreduce warp reduction mask

Open MasterJH5574 opened this issue 1 year ago • 0 comments

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.

MasterJH5574 avatar Aug 28 '24 03:08 MasterJH5574