NVIDIA / NVIDIA/cccl

[FEA]: Support indirect access to stateful predicates

Open
#10,544 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

Thrust

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

Many cub and thrust algorithms support passing functors to customize their behavior, such as the binary op for a reduction or the comparator for a sort. These functors expose a device-callable operator that is embedded into the thrust/cub algorithm. They are typically instantiated on the host and then passed directly to the call (e.g. `thrust::sort(seq.begin(), seq.end(), seq_comparator{})`), at which point thrust/cub will internally pass that parameter by value to the underlying CUDA kernel invocation, resulting in the functor being copied to the device (specifically [the `.param` memory space](https://docs.nvidia.com/cuda/parallel-thread-execution/#parameter-state-space)) and used from there. By default, the compiler will almost always try to inline the functor's call operator so that it doesn't have to introduce an extra stack frame and push more data out of registers into local memory. This has two significant potential downsides:
1. Inlining the functor everywhere has the potential to significantly balloon compile times.
2. If the functor is not stateless, since it is passed by value to the kernel every thread must maintain a copy of the functor's state in its own registers, which can consume more registers and reduce occupancy.

In the common case where the functor is stateless and not overly complex, the compile-time cost of inlining is usually worth paying to avoid the otherwise outsized negative performance impact due to the extra indirection and the potential increase in off-chip local memory traffic of not inlining the functor. For heavily stateful functors, however, the existence of the functor in the local thread state alone will increase register and local memory usage even with the call operator inlined, resulting in lowered occupancy and poor performance. In this case, with the increasing size and performance of HBM in modern GPUs the best option for both compile time _and_ runtime performance may actually be to store the comparator in global memory and pass it by reference to the kernels. That allows the compiler's live range analysis to exclude the comparator's state and focus on allocating registers to the other parts of the thrust/cub kernels, relying on the caches to service requests to the comparator state if it is frequently accessed.

### Describe the solution you'd like

Implementing this approach in thrust/cub requires some manual work to copy the functor into GPU memory and manage its lifetime relative to the kernel. It would be nice for CCCL to provide some support for automatically managing the lifetime of a stateful functor being passed through to a thrust/cub call. The ideal API might look something like:
```cpp
thrust::sort(seq.begin(), seq.end(), indirect_comparator{seq_comparator{/* complex state here */}});
```
The `indirect_comparator` type would provide RAII semantics for the device copy of the `seq_comparator` so that users could largely continue to write the same thrust/cub code that they always have and simply wrap expensive functors in these convenient wrapper functors.

I do not think there is any easy automatic way to detect when the wrapper should be used and to automatically do so due to the many lifetime concerns that could arise, nor do I think that CCCL should be in the business of trying to do so. What thrust/cub could do is warn users when they are calling functions with large functors that could have this issue. Since the compile-time warning machinery is largely done at the preprocessor level with non-portable directives like `#warning`, the best approach I can come up with for providing this warning is to abuse deprecation warnings like so:

```cpp
template
struct size_warning {};

template <>
struct [[deprecated("Object is larger than 16 bytes")]]
size_warning {};

template
void f() {
[[maybe_unused]] size_warning<(sizeof(T) >= 16)> warning;
}
```

This feature definitely isn't a must have, but it would be helpful to warn callers that they may be able to improve performance by simply wrapping their functors in something else.

### Describe alternatives you've considered

_No response_

### Additional context

@PointKernel originally discovered this issue when working with an agent on compile time optimization in cudf in https://github.com/rapidsai/cudf/pull/23448. A more complete sketch of the proposed indirection solution may be found in [this gist](https://gist.github.com/bdice/b64a753458966e237b83418b0369f93b) that @bdice wrote up after discussions in that PR.

Contributor guide

Open the contributing guide

Research direction

Start by reading the proposed indirect_comparator API and the linked design sketch, then trace how Thrust and CUB pass functors into algorithms such as thrust::sort. Define the RAII lifetime requirements for a device-resident functor and validate the wrapper against stateful comparator use; done means users can wrap expensive functors without manually managing the device copy.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
hpc, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
43/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.