NVIDIA / NVIDIA/cccl

CUB loads data from `thrust::counting_iterator` with `BLOCK_LOAD_WARP_TRANSPOSE`

Open
#660 2 comments 0 reactions 0 assignees View on GitHub
cub
Dominant language
C++
Stars
2.5k
Forks
487
Avg merge
2d 7h
Merged PRs (30d)
296

Description

### Is this a duplicate?

- [X] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)

### Area

CUB

### Is your feature request related to a problem? Please describe.

This is more of an observation rather than a feature request.

Some kernels perform BlockLoad with BLOCK_LOAD_WARP_TRANSPOSE, depending on the selected policy. When the source is not a memory location, say `thrust::counting_iterator` , loading the values directly without transpose could be more efficient.

For example, the following code runs ~10% faster on A100 when BLOCK_LOAD_DIRECT is hard-coded in AgentSelectIf (otherwise it would use BLOCK_LOAD_WARP_TRANSPOSE for this set of input parameters)
```
//nvcc -O3 -std=c++17 -arch=sm_80 --extended-lambda -Icccl/cub -Icccl/thrust -Icccl/libcudacxx/include

#include
#include
#include
#include

#include
#include

int main(){
const int N = 32*1024*1024;
thrust::device_vector d_chars(N);
thrust::sequence(d_chars.begin(), d_chars.end(), 0);

thrust::device_vector d_positions(N);
thrust::device_vector d_numSelected(1);
auto flags = thrust::make_transform_iterator(
d_chars.data().get(),
[]__host__ __device__ (char c){ return c == 'A';}
);

size_t tempsize = 0;
cub::DevicePartition::Flagged(
nullptr,
tempsize,
thrust::make_counting_iterator(0),
flags,
d_positions.data().get(),
d_numSelected.data(),
N
);

thrust::device_vector cubTemp(tempsize);

auto timeA = std::chrono::system_clock::now();
cub::DevicePartition::Flagged(
cubTemp.data().get(),
tempsize,
thrust::make_counting_iterator(0),
flags,
d_positions.data().get(),
d_numSelected.data(),
N
);
cudaDeviceSynchronize();
auto timeB = std::chrono::system_clock::now();
std::chrono::duration delta = timeB - timeA;
std::cout << delta.count() << " s\n";
}
```

### Describe the solution you'd like

Not sure. A way to manually set the load algorithm, or a way to specify that the input iterator does not perform memory accesses.

### Describe alternatives you've considered

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Research direction

Start at CUB's AgentSelectIf policy selection and the BlockLoad path, focusing on BLOCK_LOAD_WARP_TRANSPOSE versus BLOCK_LOAD_DIRECT for thrust::counting_iterator. Reproduce the provided A100 benchmark, then define and verify a supported way to select an appropriate load algorithm or recognize non-memory input without regressing ordinary loads.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.