NVIDIA / NVIDIA/cccl

thrust::tabulate is much slower than equivalent transform/for_each for wide row functor

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

Description

## Summary

`thrust::tabulate` is significantly slower than equivalent `thrust::transform`, `thrust::for_each_n`, `cub::DeviceFor::ForEach`, and a direct CUDA kernel for a row-wise functor that reads a wide row and writes one output value.

This came up while investigating libcudf row hashing performance. I minimized it to a standalone CCCL + NVBench reproducer that does not depend on cuDF or RMM.

## Reproducer

I have a self-contained reproducer with these files:

- `tabulate_reproducer.cu`
- `nvbench_main.cu`
- `CMakeLists.txt`
- `README.md`

The benchmark compares these variants over the same synthetic row-hash functor and data:

- `thrust::tabulate(output.begin(), output.end(), row_hasher)`
- `thrust::transform(counting_iterator, output.begin(), row_hasher)`
- `thrust::for_each_n(counting_iterator, lambda writes output[row])`
- `cub::DeviceFor::ForEach(counting_iterator range, lambda writes output[row])`
- direct CUDA grid-stride kernel

The core functor is:

```cpp
struct row_hasher {
cuda::std::uint64_t const* columns{};
int num_rows{};
int num_cols{};
hash_type seed{};

__device__ hash_type operator()(int row_index) const
{
auto hash = mix64_to_32(columns[row_index], seed);
for (int column_index = 1; column_index < num_cols; ++column_index) {
auto const value = columns[static_cast(column_index) * num_rows + row_index];
hash = hash_combine(hash, mix64_to_32(value, seed));
}
return hash;
}
};
```

Full reproducer source:

```cpp
// tabulate_reproducer.cu
#include

#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include

#include

namespace {

using hash_type = cuda::std::uint32_t;

__host__ __device__ hash_type hash_combine(hash_type lhs, hash_type rhs)
{
return lhs ^ (rhs + 0x9e37'79b9u + (lhs << 6) + (lhs >> 2));
}

__host__ __device__ hash_type mix64_to_32(cuda::std::uint64_t key, hash_type seed)
{
key ^= static_cast(seed) + 0x9e37'79b9'7f4a'7c15ull;
key ^= key >> 33;
key *= 0xff51afd7ed558ccdull;
key ^= key >> 33;
key *= 0xc4ceb9fe1a85ec53ull;
key ^= key >> 33;
return static_cast(key ^ (key >> 32));
}

struct row_hasher {
cuda::std::uint64_t const* columns{};
int num_rows{};
int num_cols{};
hash_type seed{};

__device__ hash_type operator()(int row_index) const
{
auto hash = mix64_to_32(columns[row_index], seed);
for (int column_index = 1; column_index < num_cols; ++column_index) {
auto const value = columns[static_cast(column_index) * num_rows + row_index];
hash = hash_combine(hash, mix64_to_32(value, seed));
}
return hash;
}
};

__global__ void direct_kernel(row_hasher hasher, hash_type* output)
{
auto row_index = static_cast(threadIdx.x + blockIdx.x * blockDim.x);
auto const stride = static_cast(blockDim.x * gridDim.x);
for (; row_index < hasher.num_rows; row_index += stride) {
output[row_index] = hasher(row_index);
}
}

void check_cuda(cudaError_t status)
{
if (status != cudaSuccess) { throw std::runtime_error(cudaGetErrorString(status)); }
}

void bench_row_hash(nvbench::state& state)
{
auto const num_rows = static_cast(state.get_int64("num_rows"));
auto const num_cols = static_cast(state.get_int64("num_cols"));
auto const variant = state.get_string("variant");

thrust::device_vector columns(
static_cast(num_rows) * num_cols);
thrust::device_vector output(num_rows);

auto const init_hasher = [data = thrust::raw_pointer_cast(columns.data()),
num_rows,
num_cols] __device__(int index) {
auto const row = index % num_rows;
auto const col = index / num_rows;
data[index] = (static_cast(col) << 32) ^
static_cast(row * 1315423911u + col);
};
thrust::for_each_n(thrust::device,
thrust::make_counting_iterator(0),
static_cast(columns.size()),
init_hasher);

row_hasher const hasher{thrust::raw_pointer_cast(columns.data()), num_rows, num_cols, 0};
auto* output_ptr = thrust::raw_pointer_cast(output.data());

state.add_global_memory_reads(static_cast(num_rows) * num_cols);
state.add_global_memory_writes(num_rows);

state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) {
auto const stream = launch.get_stream().get_stream();
if (variant == "tabulate") {
thrust::tabulate(thrust::cuda::par.on(stream), output.begin(), output.end(), hasher);
} else if (variant == "transform") {
thrust::transform(thrust::cuda::par.on(stream),
thrust::make_counting_iterator(0),
thrust::make_counting_iterator(num_rows),
output.begin(),
hasher);
} else if (variant == "for_each") {
thrust::for_each_n(thrust::cuda::par.on(stream),
thrust::make_counting_iterator(0),
num_rows,
[hasher, output_ptr] __device__(int row_index) {
output_ptr[row_index] = hasher(row_index);
});
} else if (variant == "cub_for") {
check_cuda(cub::DeviceFor::ForEach(
thrust::make_counting_iterator(0),
thrust::make_counting_iterator(num_rows),
[hasher, output_ptr] __device__(int row_index) { output_ptr[row_index] = hasher(row_index); },
stream));
} else if (variant == "direct_kernel") {
auto constexpr block_size = 256;
auto const grid_size = (num_rows + block_size - 1) / block_size;
direct_kernel<<>>(hasher, output_ptr);
check_cuda(cudaPeekAtLastError());
} else {
state.skip("unknown variant");
}
});
}

NVBENCH_BENCH(bench_row_hash)
.set_name("row_hash")
.add_int64_axis("num_rows", {1048576, 16777216})
.add_int64_axis("num_cols", {8, 32, 128, 256})
.add_string_axis("variant", {"tabulate", "transform", "for_each", "cub_for", "direct_kernel"});

} // namespace
```

```cpp
// nvbench_main.cu
#include

NVBENCH_MAIN
```

## Observed results

Hardware: NVIDIA GB10
CUDA: 13.2
CCCL: from current RAPIDS/cuDF build dependency checkout
NVBench: v0.1.0 main:v26.02.00a-1118-gf239649f8d-dirty

Command:

```bash
./build/tabulate_reproducer \
--benchmark row_hash \
--axis num_rows=1048576 \
--axis num_cols=[32,128,256] \
--axis variant=[tabulate,transform,for_each,cub_for,direct_kernel] \
--min-samples 5 \
--min-time 0.05 \
--timeout 20 \
--devices 0
```

Results:

| rows | cols | tabulate | transform | for_each | cub_for | direct_kernel |
|---:|---:|---:|---:|---:|---:|---:|
| 1,048,576 | 32 | 4.188 ms | 1.264 ms | 1.165 ms | 1.201 ms | 1.209 ms |
| 1,048,576 | 128 | 17.430 ms | 4.411 ms | 4.510 ms | 4.507 ms | 4.502 ms |
| 1,048,576 | 256 | 36.005 ms | 9.255 ms | 9.042 ms | 9.061 ms | 8.960 ms |

`thrust::tabulate` is consistently about 3.5-4x slower than the equivalent variants.

## Expected behavior

I would expect `thrust::tabulate(output.begin(), output.end(), f)` to be in the same performance family as `thrust::transform(counting_iterator, output.begin(), f)`, `thrust::for_each_n(counting_iterator, ...)`, and `cub::DeviceFor::ForEach` for this workload.

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.