NVIDIA / NVIDIA/cccl

Add CUB block reduce primitive that returns the reduced result in all threads

Open
#3,917 1 comment 0 reactions 0 assignees View on GitHub
cub
Dominant language
C++
Stars
2.5k
Forks
487
Avg merge
2d 7h
Merged PRs (30d)
296

Description

We need an `AllReduce`/`AllSum` that performs a reduction distributes the reduction across the entire block. I need this exposed in both CUB and cuda.cooperative. This pattern is fairly common in machine learning kernels. For example, you need it for softmax. Other programming frameworks have a primitive for this, such as [`MPI_Allreduce`](https://www.mpi-forum.org/docs/mpi-4.1/mpi41-report/node136.htm#Node136) in MPI.

Today, CUB block reductions only return a meaningful result (the reduced value) in the first thread of the block.

`AllReduce`/`AllSum` may need to be a new algorithm instead of a new method on `BlockReduce`, because you need a small amount of shared memory (one element's worth) to distribute the result. Existing shared memory used for the reduction itself can be reused for this purpose, although you might need to add an additional two barriers (instead of just 1) if you're reusing temporary storage, so an additional element of shared memory may be worth it. If there are `BlockReduce` specializations that use 0 shared memory, it would be particularly undesirable to increase their shared memory usage to be non-zero. Since the extra shared memory would only be needed if you're using `AllReduce`/`AllSum`, it may be better to have it in a separate algorithm.

Here's what a basic implementation would look like. If we know how the underlying BlockReduce works, we can do better; the final reduced value may already be in the temporary storage and the first barrier may be unnecessary.

```
template
__global__ T AllSum(T input)
{
using BlockReduce = cub::BlockReduce;
__shared__ union {
typename BlockReduce::TempStorage reduce;
T distribute;
} tmp;

auto sum = BlockReduce(tmp.reduce).Sum(input);

CTA_SYNC(); // May be removed if we know how the reduce algorithm works.

if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0)
tmp.distribute = sum;

CTA_SYNC(); // This is always necessary.

return tmp.distribute;
}
```

My current workaround in my cuda.cooperative examples:

```
import numpy as np
from numba import cuda
import cuda.cooperative.experimental as cudax
import math
from numba import uint8, float32
from pynvjitlink import patch
patch.patch_numba_linker(lto=True)

def cudax_block_broadcast(dtype):
context = cuda.descriptor.cuda_target.target_context
size = context.get_value_type(dtype).get_abi_size(context.target_data)

@cuda.jit
def broadcast(temp_storage, red):
cuda.syncthreads()

acc = temp_storage.view(dtype)

if cuda.threadIdx.x == 0 and cuda.threadIdx.y == 0 and cuda.threadIdx.z == 0:
acc[0] = red

cuda.syncthreads()

return acc[0]

broadcast.temp_storage_bytes = size

return broadcast

num_rows = 128
num_cols = 512

locals_per_thread = 8

threads_in_block = int(num_cols / locals_per_thread)

def maximum(a, b):
return a if a > b else b

block_load = cudax.block.load(float32, threads_in_block, locals_per_thread, algorithm='warp_transpose')
block_store = cudax.block.store(float32, threads_in_block, locals_per_thread, algorithm='warp_transpose')
block_broadcast = cudax_block_broadcast(float32)
block_max = cudax.block.reduce(float32, threads_in_block, maximum, locals_per_thread)
block_sum = cudax.block.sum(float32, threads_in_block, locals_per_thread)

temp_storage_bytes = max(block_load.temp_storage_bytes, block_store.temp_storage_bytes, block_broadcast.temp_storage_bytes, block_max.temp_storage_bytes, block_sum.temp_storage_bytes)

@cuda.jit(fastmath=True, link=block_store.files+block_load.files+block_max.files+block_sum.files)
def softmax(input, output):
temp_storage = cuda.shared.array(shape=temp_storage_bytes, dtype=uint8)
locals = cuda.local.array(shape=locals_per_thread, dtype=float32)

block_load(temp_storage, input[cuda.blockIdx.x, :], locals)

cuda.syncthreads()

max = block_broadcast(temp_storage, block_max(temp_storage, locals))

for i in range(locals_per_thread):
locals[i] = math.exp(locals[i] - max)

cuda.syncthreads()

sum = block_broadcast(temp_storage, block_sum(temp_storage, locals))

for i in range(locals_per_thread):
locals[i] = locals[i] / sum

cuda.syncthreads()

block_store(temp_storage, output[cuda.blockIdx.x, :], locals)

input = cuda.to_device(np.random.default_rng(0).random(size=(num_rows, num_cols), dtype=np.float32))
output = cuda.to_device(np.full((num_rows, num_cols), -1, dtype=np.float32))

print(input.copy_to_host())

softmax[num_rows, threads_in_block](input, output)

result = output.copy_to_host()

print(input.copy_to_host())
print(result)
print(result.sum(axis=1))

assert np.all(np.isclose(result.sum(axis=1), 1.))
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.