NVIDIA / NVIDIA/cccl

[FEA]: Provide Way To Query Shared Memory Needs At Runtime

Open
#6,771 2 comments 0 reactions 0 assignees View on GitHub
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.

Currently to use the CUB block-wide functions you must pass in a compile-time shared memory because the size must be known at compile-time:

```
using BlockReduce = cub::BlockReduce;

// Allocate shared memory for BlockReduce
__shared__ typename BlockReduce::TempStorage temp_storage;
```

This is a problem for libraries that have the size only known at runtime because they need to make occupancy and launch dimension decisions based on the shared memory usage. If it's only known after a second pass (JIT) then the shared memory requirements may not even be feasible on some architectures.

### Describe the solution you'd like

Provide a way to know how much shared memory will be used based on a runtime `size` parameter for all the block-wide functions.

### Describe alternatives you've considered

Having a third compilation pass to run a small program to get the size, then use that in JIT. This is not ideal.

### Additional context

_No response_

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.