[BUG]: Tile id computation is broken for large inputs
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
### Is this a duplicate?
- [X] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this bug and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)
### Type of Bug
Runtime Error
### Component
CUB
### Describe the bug
Some algorithms that are based on decoupled look-back compute tile index as `(blockIdx.x * gridDim.y) + blockIdx.y;`. This linearization assumes that `y` grows faster than `x`. This is not the case:
```cpp
__global__ void discover_linearization_kernel(unsigned long long int *cnt_0,
unsigned long long int *cnt_1,
unsigned long long int *cnt_2,
int max_dim_x)
{
if (threadIdx.x == 0)
{
unsigned long long int id = atomicAdd(cnt_0, 1);
if (blockIdx.y == 0)
{
if (id < max_dim_x)
{
atomicAdd(cnt_1, 1);
}
}
if (blockIdx.y == 1)
{
if (id < max_dim_x)
{
atomicAdd(cnt_2, 1);
}
}
}
}
void discover_linearization() {
unsigned long long int *cnt{};
cudaMalloc(&cnt, 3 * sizeof(unsigned long long int));
cudaMemset(cnt, 0, 3 * sizeof(unsigned long long int));
int max_dim_x{};
cudaDeviceGetAttribute(&max_dim_x, cudaDevAttrMaxGridDimX, 0);
discover_linearization_kernel<<>>(cnt, cnt + 1, cnt + 2, max_dim_x);
cudaDeviceSynchronize();
unsigned long long int cnt_h[3]{};
cudaMemcpy(cnt_h, cnt, 3 * sizeof(unsigned long long int), cudaMemcpyDeviceToHost);
printf("cnt_0: %llu\n", cnt_h[0]);
printf("cnt_1: %llu\n", cnt_h[1]);
printf("cnt_2: %llu\n", cnt_h[2]);
}
```
The code above illustrates that `x` grows faster. Wrong assumption leads to potential hangs in decoupled look-back. Decoupled look-back spins on tiles immediately to the left. Let's say that `y` dimension is equal to `4`. Block with index `[1024, 0]` will result in `tile_id = 4096`. This block will spin on results of `tile_id = 4095`, which will map to block `[1023, 1]`. Code above illustrates that this block will be scheduled after all `[*, 0]` blocks, resulting in decoupled look-back to hang (reproducer below).
This is a blocker for 64-bit problem sizes support in half of CUB. For large types, we might go as low as 64 threads with 1 item per thread. This results to lower-bound of 2^37 elements for hang to appear on modern GPUs.
Affected algorithms:
- `cub::DeviceRunLengthEncode::NonTrivialRuns`
- `cub::DeviceSelect::UniqueByKey`
- ~~`cub::DeviceSpmv`~~
- ~~`cub::DevicePartition`~~
- ~~`cub::DeviceSelect`~~
### How to Reproduce
```cpp
template
__global__ void init_kernel(ScanTileStateT tile_state, int blocks_in_grid)
{
tile_state.InitializeStatus(blocks_in_grid);
}
template
__global__ void decoupled_look_back_kernel(cub::ScanTileState tile_state)
{
using scan_op_t = cub::Sum;
using scan_tile_state_t = cub::ScanTileState;
using tile_prefix_op = cub::TilePrefixCallbackOp;
using temp_storage_t = typename tile_prefix_op::TempStorage;
// Allocate temp storage in shared memory
__shared__ temp_storage_t temp_storage;
scan_op_t scan_op{};
constexpr unsigned int threads_in_warp = 32;
const unsigned int tid = threadIdx.x;
// Hangs
const int tile_idx = (blockIdx.x * gridDim.y) + blockIdx.y;
// Works
// const int tile_idx = (blockIdx.y * gridDim.x) + blockIdx.x;
// Construct prefix op
tile_prefix_op prefix(tile_state, temp_storage, scan_op, tile_idx);
// Compute block aggregate
MessageT block_aggregate = blockIdx.x;
if (tile_idx == 0)
{
// There are no blocks to look back to, immediately set the inclusive state
if (tid == 0)
{
tile_state.SetInclusive(tile_idx, block_aggregate);
}
}
else
{
// Only the first warp in the block can perform the look back
const unsigned int warp_id = tid / threads_in_warp;
if (warp_id == 0)
{
// Perform the decoupled look-back
// Invocation of the prefix will block until the look-back is complete.
MessageT exclusive_prefix = prefix(block_aggregate);
if (tid == 0)
{
MessageT inclusive_prefix = scan_op(exclusive_prefix, block_aggregate);
}
}
}
}
template
void proof_hang() {
using scan_tile_state_t = cub::ScanTileState;
const int blocks_x = 1024 * 1024;
const int blocks_y = 2;
const auto grid = dim3(blocks_x, blocks_y);
const int blocks_in_grid = grid.x * grid.y;
// Query temporary storage requirements
std::size_t temp_storage_bytes{};
scan_tile_state_t::AllocationSize(blocks_in_grid, temp_storage_bytes);
// Allocate temporary storage
thrust::device_vector temp_storage(temp_storage_bytes);
std::uint8_t *d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
// Initialize temporary storage
scan_tile_state_t tile_status;
tile_status.Init(blocks_in_grid, d_temp_storage, temp_storage_bytes);
constexpr unsigned int threads_in_init_block = 256;
const unsigned int blocks_in_init_grid = cub::DivideAndRoundUp(blocks_in_grid,
threads_in_init_block);
init_kernel<<>>(tile_status, blocks_in_grid);
// Launch decoupled look-back
constexpr unsigned int threads_in_block = 256;
decoupled_look_back_kernel<<>>(tile_status);
// Wait for kernel to finish
cudaDeviceSynchronize();
}
```
### Expected behavior
Tile index is computed correctly, decoupled look-back based algorithms do not hang.
### Reproduction link
_No response_
### Operating System
_No response_
### nvidia-smi output
_No response_
### NVCC version
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.