NVIDIA / NVIDIA/cccl

[BUG]: Tile id computation is broken for large inputs

Open
#1,201 0 comments 0 reactions 0 assignees View on GitHub
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

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.