NVIDIA / NVIDIA/cccl

Maintain richer docstrings in cuda.cooperative

Open
#3,977 5 comments 2 reactions 1 assignee Claimed by @tpn View on GitHub
cuda.coop
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

In the past we've discussed the burden of keeping the richer docstrings in cuda.cooperative that are essentially ported over from the CUB documentation. While the information in these docstrings has been helpful, it hasn't always included all the information we want, such as the signatures of the callable algorithm that are supported. These docstrings have been a bit of a maintenance burden and have slowed down development while we're still in the prototype phase. Longer term, it would be challenging to keep these docstrings synced with CUB. Also, we know there are some architectural changes coming that will simplify the cuda.cooperative API, so work spent on these docstrings may be wasted. Ideally in the future, the documentation can be automatically generated from the CUB documentation. For now, for some algorithms, like block scan, we've decided to go with much simpler and lightweight 1-line docstrings.

Despite all the reasons I've just enumerated, I think it's a mistake to ditch the richer docstrings.

If the feature isn't documented no one will use it. If we don't document it **as we develop it**, it will become much more work to document it in the future, and we may end up forgoing documentation work in favor of more exciting development work. We've been burned on the past by not documenting experimental work, such as the Thrust async interfaces, which were never used because they were not documented.

We should keep the richer docstrings around, and improve them, even if it's work we'll throw away.

Here are the requirements I have in mind for cuda.cooperative docstrings for functions that instantiate an algorithm:
- One to two sentences explaining that the function produces a callable and what that callable does.
- Description of parameters to the function that creates the callable.
- List documenting the signatures supported by the callable.
- Important warnings/caveats.
- Code example that is tested as part of CI.
- Link to the corresponding CUB docs.

[Example of this, which I consider best practice - `block.load` docs from my recent PR](https://github.com/NVIDIA/cccl/blob/e7af76f4a4960d2c15195016aee96fdf0248d76d/python/cuda_cooperative/cuda/cooperative/experimental/block/_block_load_store.py#L42-L83) #3161:

```
"""Creates an operation that performs a block-wide load.

Returns a callable object that can be linked to and invoked from device code. It can be
invoked with the following signatures:

- `(src: numba.types.Array, dest: numba.types.Array) -> None`: Each thread loads
`items_per_thread` items from `src` into `dest`. `dest` must contain at least
`items_per_thread` items.

Different data movement strategies can be selected via the `algorithm` parameter:

- `algorithm="direct"` (default): A blocked arrangement of data is read directly from memory.
- `algorithm="striped"`: A striped arrangement of data is read directly from memory.
- `algorithm="vectorize"`: A blocked arrangement of data is read directly from memory using CUDA's built-in vectorized loads as a coalescing optimization.
- `algorithm="transpose"`: A striped arrangement of data is read directly from memory and is then locally transposed into a blocked arrangement.
- `algorithm="warp_transpose"`: A warp-striped arrangement of data is read directly from memory and is then locally transposed into a blocked arrangement.
- `algorithm="warp_transpose_timesliced"`: A warp-striped arrangement of data is read directly from memory and is then locally transposed into a blocked arrangement one warp at a time.

For more details, [read the corresponding CUB C++ documentation](https://nvidia.github.io/cccl/cub/api/classcub_1_1BlockLoad.html).

Args:
dtype: Data type being loaded
threads_per_block: The number of threads in a block, either an integer or a tuple of 2 or 3 integers
items_per_thread: The number of items each thread loads
algorithm: The data movement algorithm to use

Example:
The code snippet below illustrates a striped load and store of 128 integer items by 32 threads, with
each thread handling 4 integers.

.. literalinclude:: ../../python/cuda_cooperative/tests/test_block_load_store_api.py
:language: python
:dedent:
:start-after: example-begin imports
:end-before: example-end imports

.. literalinclude:: ../../python/cuda_cooperative/tests/test_block_load_store_api.py
:language: python
:dedent:
:start-after: example-begin load_store
:end-before: example-end load_store
"""
```

[Example of what I don't want - status quo of `block.inclusive_sum`](https://github.com/NVIDIA/cccl/blob/main/python/cuda_cooperative/cuda/cooperative/experimental/block/_block_scan.py#L209-L211):

```
"""
Computes an inclusive block-wide prefix sum.
"""
```

This doesn't tell the user that it's a function that /instantiates/ an algorithm which you then have to call, doesn't explain any of the parameters, etc.

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.