[cuda.compute]: add support for bfloat16 - #10933
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesBfloat16 support is added to CCCL C++ and Python compute APIs. Type conversion, JIT generation, interop, reductions, scans, sorting, serialization, histogram handling, and CI tests now handle Bfloat16 support
Assessment against linked issues
Suggested reviewers: Merge Risk: 🟡 Moderate · up to At the current head, bfloat16 support has a known failing histogram test pending PR Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
c/parallel/test/test_reduce.cpp (1)
525-525: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: Qualify
operation_t::operator cccl_op_t()asconst, then declareopconst. Aconst operation_tcannot use the current conversion inreduce(...).Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 91e893a9-45c6-44ef-b092-36ca3460efd6
📒 Files selected for processing (23)
c/parallel.v2/include/cccl/c/types.hc/parallel.v2/src/hostjit/codegen/cub_call.cppc/parallel.v2/src/hostjit/codegen/types.cppc/parallel/include/cccl/c/types.hc/parallel/src/histogram.cuc/parallel/src/jit_templates/template_pre.h.inc/parallel/src/kernels/operators.cppc/parallel/src/util/serialization.hc/parallel/src/util/types.hc/parallel/test/test_reduce.cppc/parallel/test/test_util.hci/test_cuda_compute_python.shci/windows/test_cuda_compute_python.ps1python/cuda_cccl/cuda/compute/_bindings.pyipython/cuda_cccl/cuda/compute/_bindings_impl.pyxpython/cuda_cccl/cuda/compute/_cccl_interop.pypython/cuda_cccl/cuda/compute/_cpp_compile.pypython/cuda_cccl/cuda/compute/_jit.pypython/cuda_cccl/cuda/compute/_utils/protocols.pypython/cuda_cccl/cuda/compute/iterators/_common.pypython/cuda_cccl/cuda/compute/types.pypython/cuda_cccl/pyproject.tomlpython/cuda_cccl/tests/compute/test_bfloat16.py
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| return _register_struct_with_numba(struct_class) | ||
|
|
||
| # Numba has no bfloat16 support; fail with a clear message rather than a | ||
| # cryptic numba error deep in the compilation pipeline. |
This comment has been minimized.
This comment has been minimized.
|
CI won't pass until the histogram bfloat16 bug fix is merged #10954 |
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e7d1550d-b362-4d7b-b80d-12cdaaa9ce82
📒 Files selected for processing (4)
c/parallel/src/histogram.cuc/parallel/test/test_util.hpython/cuda_cccl/cuda/compute/types.pypython/cuda_cccl/tests/compute/test_bfloat16.py
Included review availability: Your plan provides up to 12 included reviews per hour; 8 remain after this review.
| def test_histogram_even(): | ||
| num_samples = 1000 | ||
| num_bins = 8 | ||
| h_samples = random_bfloat16(num_samples, low=0.0, high=8.0, seed=29) | ||
| d_samples = DeviceArray.from_numpy(h_samples) | ||
| d_histogram = DeviceArray.from_numpy(np.zeros(num_bins, dtype=np.uint32)) | ||
|
|
||
| cuda.compute.histogram_even( | ||
| d_samples=d_samples, | ||
| d_histogram=d_histogram, | ||
| num_output_levels=num_bins + 1, | ||
| lower_level=np.array([0.0], dtype=BFLOAT16)[0], | ||
| upper_level=np.array([8.0], dtype=BFLOAT16)[0], | ||
| num_samples=num_samples, | ||
| ) | ||
|
|
||
| # histogram_even uses half-open [lower, upper) bins; samples that rounded | ||
| # up to exactly 8.0 in bfloat16 are out of range. | ||
| samples = h_samples.astype(np.float64) | ||
| expected = np.histogram(samples[samples < 8.0], bins=num_bins, range=(0, 8))[0] | ||
| np.testing.assert_array_equal(d_histogram.copy_to_host(), expected) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
important: Include the histogram bfloat16 fix before this test runs in required CI. The PR objective states that this test cannot pass until PR #10954 merges. Rebase or stack this change on that fix, then run the affected test. As per coding guidelines, “Build and run relevant tests before submitting contributions to ensure existing functionality is not broken.”
Source: Coding guidelines
This comment has been minimized.
This comment has been minimized.
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8c59129e-e356-47bf-96f9-fc30e2e6e584
📒 Files selected for processing (2)
python/cuda_cccl/cuda/compute/algorithms/_histogram.pypython/cuda_cccl/tests/compute/test_bfloat16.py
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.
| @pytest.mark.xfail( | ||
| raises=NotImplementedError, | ||
| reason="CUB's HistogramEven bins bfloat16 samples incorrectly, so " | ||
| "make_histogram_even rejects bfloat16; " | ||
| "see https://github.com/NVIDIA/cccl/issues/10940", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
while IFS= read -r file; do
rg -n -C 3 'xfail_strict|strict\s*=' "$file" || true
done < <(fd -a -t f '^(pyproject\.toml|pytest\.ini|setup\.cfg|tox\.ini)$' .)
rg -n -C 8 '`@pytest`\.mark\.xfail|histogram' \
python/cuda_cccl/tests/compute/test_bfloat16.pyRepository: NVIDIA/cccl
Length of output: 201
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- pytest configuration files ---'
find . -maxdepth 3 -type f \( -name 'pyproject.toml' -o -name 'pytest.ini' -o -name 'setup.cfg' -o -name 'tox.ini' \) -print
printf '%s\n' '--- strict-xfail settings ---'
rg -n -C 4 'xfail_strict|strict[[:space:]]*=' . --glob 'pyproject.toml' --glob 'pytest.ini' --glob 'setup.cfg' --glob 'tox.ini' || true
printf '%s\n' '--- test marker and body ---'
sed -n '470,525p' python/cuda_cccl/tests/compute/test_bfloat16.py
printf '%s\n' '--- test guidance and required project docs ---'
for file in .agent/skills/cccl-test/SKILL.md CONTRIBUTING.md docs/infrastructure/ci/references/ci_overview.rst; do
if [ -f "$file" ]; then
printf '%s\n' "--- $file ---"
sed -n '1,220p' "$file"
else
printf '%s\n' "MISSING: $file"
fi
doneRepository: NVIDIA/cccl
Length of output: 23523
important: Add strict=True to this pytest.mark.xfail marker. The repository has no global strict-xfail setting. Without strict mode, an XPASS will not fail CI.
🥳 CI Workflow Results🟩 Finished in 1h 42m: Pass: 100%/531 | Total: 5d 19h | Max: 1h 17m | Hits: 100%/624069See results here. |
3 participants
Description
closes #10906
This PR allows using
bfloat16in cuda.compute, which corresponds to_nv_bfloat16in CUDA C++. One wrinkle is that numpy does not supportbfloat16, so we have to rely on theml_dtypespackage. I decided to not add this to the list of required dependencies, rather users who want to usebfloat16must installml_dtypesthemselves, which is how CuPy does it.I verified that this is compatible with CuPy arrays and PyTorch tensors.
Checklist