You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
cuda.core: map ml_dtypes bfloat16 to kDLBfloat in DLPack export - #3359
_smv_dtype_numpy_to_dlpack picked the DLPack type code from numpy.dtype.kind, but ml_dtypes.bfloat16 has kind "V", so a bf16 view fell through to the unsupported-dtype BufferError and could not be exported — even though the import side already accepts kDLBfloat (mapping it back to bfloat16 when ml_dtypes is installed) and _tensor_map.pyx maps bfloat16 the same way. A bf16 tensor from PyTorch/CuPy/JAX can be viewed but not re-exported.
This adds an explicit bfloat16 branch ahead of the kind chain that emits kDLBfloat with 16 bits, guarded by the module-level ml_dtypes availability check already used on the import side.
The new test round-trips through StridedMemoryView.from_dlpack itself, so it cannot skip on stacks where NumPy's own DLPack support for extended dtypes is unavailable — the hole that let the existing test_dlpack_export_roundtrip_dtypes[bfloat16] miss this via its NumPy-consumer probe.
Evidence from the released cuda-core 1.2.2 wheel (host memory; numpy 2.3.5, ml_dtypes 0.6.0), matching the issue's reproducer:
>>> view.__dlpack__()
BufferError: Unsupported dtype for DLPack export: bfloat16
with float16/uint16 views of the same buffer exporting fine. Additionally: cython -3 on cuda/core/_memoryview.pyx produces identical diagnostics on base and on this branch (the two existing cuStreamWaitEvent assign notes are pre-existing; none added), and the new branch's dtype equality check matches bfloat16 while not matching bool/uint16/float16/int8. Full build and suite validation runs in this repository's CI.
The documentation is up to date with these changes.
AI assistance disclosure: the change and test were prepared with an AI coding agent (GLM-5.3) under my direction; I reviewed each line, verified the reproducer and evidence above, and stand behind the change. The new test carries @pytest.mark.agent_authored(model="glm-5.3") per the provenance rules in AGENTS.md.
It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.
No actionable comments were generated in the recent review.
Summary
Summary by CodeRabbit
Bug Fixes
DLPack exports now recognize bfloat16 arrays, improving compatibility when sharing them with other tools. Importing an exported array preserves its data type, shape, and data pointer, supporting reliable round trips. Other unsupported data types continue to raise an error. The bfloat16 round-trip check is skipped when optional type support is unavailable.
DLPack export now maps ml_dtypes.bfloat16 to kDLBfloat. A round-trip test checks the imported dtype, shape, and data pointer.
Changes: bfloat16 DLPack export
Layer / File(s)
Summary
Map and round-trip bfloat16 cuda_core/cuda/core/_memoryview.pyx, cuda_core/tests/test_utils_dlpack.py
_smv_dtype_numpy_to_dlpack maps ml_dtypes.bfloat16 to kDLBfloat. The test checks the imported dtype, shape, and data pointer. It skips when ml_dtypes is unavailable.
Assessment against linked issues
Objective
Addressed
Explanation
Enable StridedMemoryView bfloat16 DLPack export and same-memory re-import [#3337]
✅
Suggested reviewers:rluo8
Change: Bug fix · Severity of issue fixed: Medium
Merge Risk: Low · up to 2b3a0
The bfloat16 regression test can silently skip if its dependency is missing. The change is otherwise mergeable with that coverage limitation noted.
- [ ] Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Tip
🎯 Customize your risk assessments
Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.
✅ Automatic approvals, with your rules
Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.
Reviewing files that changed from the base of the PR and between 6b8430a and 2b3a074.
Files selected for processing: 1
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
Reviewing files that changed from the base of the PR and between 4010017 and eb3fdb0.
Files selected for processing: 2
cuda_core/cuda/core/_memoryview.pyx
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Addressed: test_dlpack_export_bfloat16_roundtrip now requests init_cuda, since .view() reaches the driver via cuPointerGetAttributes — thanks for catching that this also matches the other driver-touching tests in this file.
Reviewing files that changed from the base of the PR and between eb3fdb0 and 3d46b39.
Files selected for processing: 1
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Thanks, LGTM in general, but I'd ask you to move the branch and leave a bit more concrete bread-crumb that this test really should be irrelevant soon (and why).
The reason will be displayed to describe this comment to others. Learn more.
Can you move this to be the last branch here, and also can you make the check np_dtype.type is bfloat16 (we already checked byte-order, so that is fine).
Thanks @seberg — both addressed: the bfloat16 branch is now the last one before the error, using np_dtype.type is bfloat16 (safe when ml_dtypes is absent, since no dtype type is None), and the test docstring carries the bread-crumb that it becomes unnecessary once _assert_dlpack_export_roundtrip tests bfloat16 (ml_dtypes>0.6.0).
The reason will be displayed to describe this comment to others. Learn more.
Nitpick comments: 1
cuda_core/tests/test_utils_dlpack.py: 1
88-93: Trivial · Data Integrity & Integration · Quick win
Assert the exported DLPack bit width in the bfloat16 round-trip test.
The bfloat16 round-trip test checks the imported dtype, shape, and pointer. It does not check the exported bits value. The importer maps kDLBfloat to bfloat16 without checking bits, so a regression that exports kDLBfloat with a wrong bit width would not fail this test.
The helpers this test file defines are sufficient to inspect the capsule. Use _PyCapsule_GetPointer with the versioned capsule name and cast the result to _DLManagedTensorVersioned. Then assert the code and the bit width. Export the capsule once and keep the reference alive while the assertions run.
Reviewing files that changed from the base of the PR and between a876815 and 6b8430a.
Files selected for processing: 2
cuda_core/cuda/core/_memoryview.pyx
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
_smv_dtype_numpy_to_dlpack dispatched on numpy dtype kind, but
ml_dtypes.bfloat16 has kind "V", so bf16 views fell through to the
unsupported-dtype BufferError and could not be re-exported, even though
the DLPack import path already accepts kDLBfloat and _tensor_map.pyx
maps bfloat16. Detect bfloat16 up front and export it as kDLBfloat
(16 bits).
The new round-trip test goes through StridedMemoryView.from_dlpack so
it cannot skip on stacks where NumPy's own DLPack support for extended
dtypes is unavailable, unlike the existing parametrized round-trip
test's NumPy-consumer probe.
Signed-off-by: Sainath-Reddy7 <180950590+Sainath-Reddy7@users.noreply.github.com>
Also took CodeRabbit's quick-win nitpick: the bfloat16 round-trip test now pins the exported dtype.code == 4 (kDLBfloat in dlpack.h) and bits == 16 by inspecting the unversioned capsule with the file's existing ctypes helpers.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
cuda.coreEverything related to the cuda.core module
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
seberg
left a comment
Description
closes #3337
_smv_dtype_numpy_to_dlpackpicked the DLPack type code fromnumpy.dtype.kind, butml_dtypes.bfloat16has kind"V", so a bf16 view fell through to the unsupported-dtypeBufferErrorand could not be exported — even though the import side already acceptskDLBfloat(mapping it back tobfloat16whenml_dtypesis installed) and_tensor_map.pyxmapsbfloat16the same way. A bf16 tensor from PyTorch/CuPy/JAX can be viewed but not re-exported.This adds an explicit bfloat16 branch ahead of the kind chain that emits
kDLBfloatwith 16 bits, guarded by the module-levelml_dtypesavailability check already used on the import side.The new test round-trips through
StridedMemoryView.from_dlpackitself, so it cannot skip on stacks where NumPy's own DLPack support for extended dtypes is unavailable — the hole that let the existingtest_dlpack_export_roundtrip_dtypes[bfloat16]miss this via its NumPy-consumer probe.Evidence from the released cuda-core 1.2.2 wheel (host memory; numpy 2.3.5, ml_dtypes 0.6.0), matching the issue's reproducer:
with float16/uint16 views of the same buffer exporting fine. Additionally:
cython -3oncuda/core/_memoryview.pyxproduces identical diagnostics on base and on this branch (the two existing cuStreamWaitEvent assign notes are pre-existing; none added), and the new branch's dtype equality check matchesbfloat16while not matching bool/uint16/float16/int8. Full build and suite validation runs in this repository's CI.Checklist
AI assistance disclosure: the change and test were prepared with an AI coding agent (GLM-5.3) under my direction; I reviewed each line, verified the reproducer and evidence above, and stand behind the change. The new test carries
@pytest.mark.agent_authored(model="glm-5.3")per the provenance rules in AGENTS.md.