[CUCO][FEA]: Use plain CUDA atomics instead of cuda::atomic_ref
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
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
CUDA Experimental (cudax)
### Is your feature request related to a problem? Please describe.
The current cudax cuco implementations use `cuda::atomic_ref` for HLL and open-addressing atomic operations. There are several practical issues with this.
From what we have tested so far, `cuda::atomic_ref` does not always deliver the same codegen as plain CUDA atomics. In [NVIDIA/cuCollections#808](https://github.com/NVIDIA/cuCollections/pull/808#issuecomment-4784118925), enabling `cuda::atomic_ref::fetch_or` instead of `atomicOr` introduced a 3.44% average add throughput regression on RTX PRO 6000 Blackwell Max-Q with CUDA 13.1, and about 3.6% for filter sizes up to 128 MiB. This is small but not negligible for hot cuco paths. The original discussion also confirmed [suboptimal codegen from `atomic_ref::fetch_or`](https://github.com/NVIDIA/cuCollections/pull/808#discussion_r3463275968).
The relevant libcu++ device atomic implementations currently use volatile inline PTX with a `"memory"` clobber. This gives the compiler front end less visibility across the atomic and can constrain optimization and instruction scheduling compared with compiler-recognized CUDA atomic intrinsics. We have not isolated this as the sole reason for the measured regression, but the ablation above confirms the codegen and performance difference.
There are also correctness issues:
- [NVIDIA/cuCollections#816](https://github.com/NVIDIA/cuCollections/pull/816) needs to define the CCCL-internal `_CCCL_ATOMIC_UNSAFE_AUTOMATIC_STORAGE` workaround for CUDA below 13.1. This avoids an NVCC miscompile in `cuda::atomic_ref` automatic-storage dispatch that can cause open-addressing inserts to be lost, as reported in [NVIDIA/cuCollections#804](https://github.com/NVIDIA/cuCollections/issues/804).
We cannot rely on CCCL enabling this workaround globally. The macro deliberately disables valid automatic-storage handling for the whole translation unit, while cuco only needs these atomics for its global/shared storage.
- We also encountered a downstream Blackwell case in [rapidsai/cudf#23229](https://github.com/rapidsai/cudf/pull/23229). With CUDA 13.0, selecting the existing fallback using two 64-bit `cuda::atomic_ref` CAS loops still produced incorrect `DECIMAL128` groupby sums. Replacing those loops with native 64-bit `atomicAdd` operations made the reproducer pass.
We do not yet fully understand the root cause, and this is not an apples-to-apples comparison of equivalent CAS implementations, but it is another case where using plain CUDA atomics avoided the issue.
cuco previously had private atomic helpers that dispatched by type width and thread scope to `atomicCAS`, `atomicExch`, and their block/system variants. HLL similarly dispatched to `atomicMax_block`, `atomicMax`, or `atomicMax_system`. These implementations were replaced with `cuda::atomic_ref` in [NVIDIA/cuCollections#531](https://github.com/NVIDIA/cuCollections/pull/531) after the earlier [CCCL #1008](https://github.com/NVIDIA/cccl/issues/1008) regression was fixed.
The performance comparison in #531 showed parity on T4 at that time. The newer #808 result is a separate workload and toolchain observation and should not be confused with the resolved #1008 issue.
### Describe the solution you'd like
Introduce an internal cudax cuco atomic utility that dispatches by operation, supported type/width, and thread scope to plain CUDA atomic intrinsics.
This utility should:
- Replace the current `cuda::atomic_ref` usages in HLL and open addressing, including CAS, store/exchange, add, max, and relaxed atomic loads.
- Preserve all currently supported types and thread scopes. Small types, 128-bit types, and cluster scope should be handled explicitly instead of silently losing support.
- Use a narrow fallback only where no suitable plain CUDA atomic operation exists.
- Add correctness coverage for the packed-CAS paths and Release builds with CUDA below 13.1.
- Compare HLL and fixed-capacity-map performance against the current `atomic_ref` implementation on representative GPUs and CUDA versions.
This is not a literal revert of #531 because the current cudax implementation has more operations and supported configurations than the old helpers covered.
This serves as an internal temporary workaround. Once the above codegen and compiler issues have been resolved and `cuda::atomic_ref` provides comparable performance, we can replace these helpers with `cuda::atomic_ref` again.
### Describe alternatives you've considered
Keep using `cuda::atomic_ref` and maintain toolchain-specific workarounds.
This would not address the measured codegen regression, and the automatic-storage workaround cannot be safely enabled globally for all CCCL users. Maintaining a small cuco-internal dispatch layer is more reliable while these issues remain unresolved.
### Additional context
Related history:
- [NVIDIA/cuCollections#469](https://github.com/NVIDIA/cuCollections/issues/469)
- [NVIDIA/cuCollections#531](https://github.com/NVIDIA/cuCollections/pull/531)
- [NVIDIA/cuCollections#808](https://github.com/NVIDIA/cuCollections/pull/808)
- [NVIDIA/cuCollections#816](https://github.com/NVIDIA/cuCollections/pull/816)
- [rapidsai/cudf#23229](https://github.com/rapidsai/cudf/pull/23229)
Contributor guide
Research direction
Start by locating the cudax cuco HLL and open-addressing implementations that use cuda::atomic_ref, then compare them with the previously mentioned private atomic helpers. Define the dispatch needed for each operation, supported width, type, and thread scope before changing behavior. Done means plain CUDA atomics cover the current configurations, correctness tests include packed-CAS and pre-13.1 Release builds, and HLL and fixed-capacity-map benchmarks are compared.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- hpc, performance, testing-qa
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100