[TIR] Allreduce broadcast result to each thread in multi-warp case

[MasterJH5574]

PR #15327 introduces the warp-level primitive support in multi-warp allreduce. However, due to the specialty of the two-stage shuffle-down reduction implementation of the allreduce in multi-warp scenarios, PR #15327 did not broadcast the allreduce result to each reduction thread. This behavior does not align with the semantics of allreduce and is not ideal for many use cases. Therefore, this PR completes the implementation by inserting a stage of writing the reduction results to shared memory, so that each reduction thread across all the reduction warps can access the reduction results.

This shared memory write-back stage will only be inserted in multi-warp allreduce cases. In single-warp allreduce, a shfl_sync is used to broadcast the reduction results across reduction threads. Since in multi-warp settings we cannot leverage warp-level primitives to broadcast the value, we can only make use of shared memory.

The numerical correctness are verified locally.

[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.

• cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[MasterJH5574]

cc @yzh119 @tqchen

[junrushao]

LGTM! I was curious if there's any performance implication of this change?

[MasterJH5574]

LGTM! I was curious if there's any performance implication of this change?

@junrushao I didn’t measure. For platform like CUDA with warp size 32, the additional shared memory will have at most 16 elements, and I assume this overhead is negligible. Nevertheless, in multi-warp reduction settings, the current impl with warp-level primitive leverage will be at least no slower than the status before #15327, which allocated large shared memory and used naive allreduce implementation over shared memory.

On the other hand, to fulfill the semantics of allreduce, we have to compromise the shared memory here.