NVIDIA / NVIDIA/cudf

[FEA] Migrate libcudf from Thrust to CUB-based algorithms

Open
#21,272 1 comment 2 reactions 0 assignees View on GitHub
feature request libcudf
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

## Motivation

libcudf should incrementally migrate from Thrust algorithms toward CUB-based internal implementations. This migration provides several key benefits:

### 1. Native 64-bit offset support
CUB has broad support for 64-bit index types (https://github.com/NVIDIA/cccl/issues/50), which is critical for:
- Supporting large data ranges without overflow concerns
- Smoothing _potential_ future changes in `size_type` from `int32_t` to `int64_t`

Thrust's dynamic dispatch between 32-bit and 64-bit offsets requires compiling two kernels, which is expensive for compile time and binary size. We currently force all Thrust kernels to use 32-bit offsets.

### 2. Elimination of unnecessary stream synchronizations
Even with nosync everywhere (https://github.com/rapidsai/cudf/pull/20807, https://github.com/rapidsai/cudf/pull/20978), Thrust algorithms often require stream synchronization to return scalar results to the host via pageable memory. Custom CUB-based implementations using pinned memory can avoid these synchronizations:
- Pinned memory D2H transfers don't block other GPU work
- Critical for multithreaded readers (Parquet/ORC) where one thread's sync can stall another's Copy Engine usage
- Enables truly async operations when return values aren't needed

### 3. Better performance diagnostics
Performance regressions in CUB are easier to diagnose than in Thrust, because Thrust typically calls CUB. By removing a layer of indirection between libcudf and the underlying implementation, we have more direct control over algorithm behavior.

### 4. Reduced complexity
Historically one reason to prefer Thrust was its brevity and STL-like interface. This is less significant as CCCL adds memory resource control to CUB via env-based overloads (https://github.com/NVIDIA/cccl/issues/5635). Single-dispatch CUB calls eliminate the need to call twice (once with `nullptr` to get size, again with buffer pointer).

## Completed Work

The following PRs have established the pattern and migrated several algorithms:

### `copy_if`
- PR #20861: Replaced `thrust::copy_if` and `thrust::count_if` with CUB + pinned memory wrappers
- PR #21156: Added `cudf::detail::copy_if` with stencil using `cub::DeviceSelect::FlaggedIf` for 64-bit range support
- PR #21268: Replaced remaining `thrust::copy_if` calls with `cudf::detail::copy_if`

### `reduce_by_key`
- PR #20860: Replaced `thrust::reduce_by_key` with CUB + pinned memory wrapper
- PR #21184: Added `cudf::detail::reduce_by_key_async` for cases where return iterator pair is not needed

### `all_of`, `any_of`
- PR #20822: Replaced `thrust::all_of`, `thrust::any_of` with CUB-based implementations using pinned memory

## Remaining Work

Identify and migrate other Thrust algorithms that:
- Return scalar values, and CUB could do so with pinned memory (reducing transfer overhead and avoiding a sync from D2H copy)
- Are used in hot paths or multithreaded contexts (e.g. I/O paths that are especially performance-sensitive and must avoid syncs)
- Would benefit from 64-bit offset support

As noted in PR #21268 review comments, we may want to break out algorithms into separate `.cuh` files as the collection grows. Current location: `cudf/detail/utilities/algorithm.cuh`

## Implementation Notes

The established pattern for CUB-based wrappers:
1. Use `rmm::device_uvector` or `cudf::detail::device_scalar` for device-side output
2. Use pinned memory bounce buffer for D2H transfer (via `cudf::detail::device_scalar` transparently)
3. Provide `*_async` variants when callers don't need the return value
4. Support 64-bit ranges via `cuda::std::size_t` / `cuda::std::distance`
5. Wrap CUB calls with `CUDF_CUDA_TRY` for error handling

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.