NVIDIA / NVIDIA/cccl

cub block reductions fail to compile correctly with nvrtc for certain block sizes

Open
#856 6 comments 0 reactions 1 assignee Claimed by @jrhemstad View on GitHub
cub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

Using cub's block reductions in kernels compiled using nvrtc (using Jitify), fail to compile for specific block sizes. See the error produced below, where template type deduction seems to be failing in thread_store.cuh. Curiously this only occurs for "unusual" block sizes, and more regular block sizes compile fine.

```c++
../../thread/thread_store.cuh(351): error: argument list for template "cub::IterateThreadStore::Dereference [with COUNT=0, MAX=VOLATILE_MULTIPLE]" is missing
detected during:
instantiation of "void cub::ThreadStore(T *, T, cub::Int2Type<5>, cub::Int2Type<1>) [with T=quda::complex]"
(410): here
instantiation of "void cub::ThreadStore(OutputIteratorT, T) [with MODIFIER=cub::STORE_VOLATILE, OutputIteratorT=quda::complex *, T=quda::complex]"
specializations/warp_reduce_smem.cuh(149): here
instantiation of "T cub::WarpReduceSmem::ReduceStep(T, int, ReductionOp, cub::Int2Type) [with T=quda::complex, LOGICAL_WARP_THREADS=18, PTX_ARCH=600, ALL_LANES_VALID=true, FOLDED_ITEMS_PER_LANE=1, ReductionOp=cub::Sum, STEP=0]"
specializations/warp_reduce_smem.cuh(349): here
instantiation of "T cub::WarpReduceSmem::Reduce(T, int, ReductionOp) [with T=quda::complex, LOGICAL_WARP_THREADS=18, PTX_ARCH=600, ALL_LANES_VALID=true, FOLDED_ITEMS_PER_LANE=1, ReductionOp=cub::Sum]"
specializations/block_reduce_warp_reductions.cuh(184): here
instantiation of "T cub::BlockReduceWarpReductions::Sum(T, int) [with T=quda::complex, BLOCK_DIM_X=9, BLOCK_DIM_Y=2, BLOCK_DIM_Z=1, PTX_ARCH=600, FULL_TILE=true]"
cub/block/block_reduce.cuh(500): here
instantiation of "T cub::BlockReduce::Sum(T) [with T=quda::complex, BLOCK_DIM_X=9, ALGORITHM=cub::BLOCK_REDUCE_WARP_REDUCTIONS, BLOCK_DIM_Y=2, BLOCK_DIM_Z=1, PTX_ARCH=600]"
```

A simple patch to fix the above problem is to force the template type deduction, e.g.
```c++
--- a/include/externals/cub/thread/thread_store.cuh
+++ b/include/externals/cub/thread/thread_store.cuh
@@ -348,7 +348,7 @@
__device__ __forceinline__ void ThreadStoreVolatilePtr(
for (int i = 0; i < SHUFFLE_MULTIPLE; ++i)
reinterpret_cast(words)[i] = reinterpret_cast(&val)[i];

- IterateThreadStore<0, VOLATILE_MULTIPLE>::template Dereference(
+ IterateThreadStore<0, VOLATILE_MULTIPLE>::template Dereference(
reinterpret_cast(ptr),
words);
}
```
fixes the compilation problem and the kernels executes correctly.

Whose fault this is, I don't know. It looks like this issue may be hit elsewhere in cub, but my use case doesn't extend much beyond reductions at present.

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.