CUB loads data from `thrust::counting_iterator` with `BLOCK_LOAD_WARP_TRANSPOSE`
- 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
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