CUDA: Fix data-races when reusing SMEM in block_reduce (#26385)
* CUDA: Fix data-races when reusing block_reduce block_reduce currently doesn't resync after reading from SMEM, causing potential data-races when reusing SMEM for multiple reductions. One may consider simply always adding this in block_reduce, but this comes at a potential perf cost * double-buffering for single-row softmax * double-buffering for norm as well * Add comment * Add explanatory comment to block_reduce * Specify need for + do memory barrier only in multi-warp scenario * Implement review-suggestion from @gaugarg-nv
This commit is contained in:
@@ -627,7 +627,8 @@ template <typename T> struct block_reduce_policy<block_reduce_method::MAX, T> {
|
||||
};
|
||||
|
||||
template <block_reduce_method reduce_method_t, const unsigned int block_size_template = 0, typename T>
|
||||
static __device__ T block_reduce(T val, T * shared_vals) {
|
||||
static __device__ T block_reduce(T val, [[maybe_unused]] T * shared_vals) {
|
||||
// for multi-warp reductions, callers must not reuse shared_vals until all reads from this invocation have completed
|
||||
val = block_reduce_policy<reduce_method_t, T>::reduce(val);
|
||||
const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
|
||||
if (block_size > WARP_SIZE) {
|
||||
|
||||
Reference in New Issue
Block a user