[cuda.compute]: reject RawOp values with mismatched complex layout - #11525
saiyambharara wants to merge 1 commit into
Conversation
saiyambharara
changed the title
NumPy's align=True packer aligns complex64 and complex128 to their real component instead of libcu++'s alignas(2 * sizeof(T)). A RawOp casting to a native struct can silently read or write the wrong bytes when a complex field lands at a different offset (NVIDIA#11347). RawOp.compile() now checks input and output layouts against native CUDA/C++, including nested structs, subarrays, and pointers, and raises a clear error instead of building a broken kernel. Python callable operators are not affected.
ab96be6 to
2bdbfe2
Compare
📝 SummarySummary by CodeRabbit
WalkthroughRawOp compilation now validates CUDA type-descriptor layouts against native CUDA/C++ layouts. The validation covers complex types, nested structs, pointers, subarrays, offsets, size, and alignment. Regression tests cover rejection, compatibility, and end-to-end execution. ChangesRawOp layout validation
Assessment against linked issues
Suggested reviewers: Priority: ➖ Normal Change: Bug fix · Severity of issue fixed: Medium Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cuda_cccl/cuda/compute/op.py (1)
68-71: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: checking a pointee as a root rejects valid programs when the pointee is a scalar. Example: a struct field of type
complex64*._check_raw_op_layoutthen reaches the bare-type branch and compares NumPy's 4-byte alignment with the native 8, so compilation fails. No byte layout differs in that case: the element size is identical, and the pointed-to allocation is not laid out bycuda.compute. Restrict the pointee check to struct pointees, where field offsets and stride can actually disagree.if isinstance(td, cccl_types.PointerTypeDescriptor): # The pointee is a separate allocation: check it as its own root. - _check_raw_op_layout(td.pointee, path=f"*{path}") + # Only its internal layout can disagree; a scalar pointee's alignment + # is a property of the allocation, not of the layout. + if isinstance(td.pointee, cccl_types.StructTypeDescriptor): + _check_raw_op_layout(td.pointee, path=f"*{path}") return
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository: NVIDIA/cccl/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7a3eb549-236e-444b-9e09-ed30d8925c9a
📒 Files selected for processing (2)
python/cuda_cccl/cuda/compute/op.pypython/cuda_cccl/tests/compute/test_raw_op_complex_layout.py
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
1 participant
Description
closes #11347
cuda.compute builds complex64/complex128 layout in structs and subarrays through NumPy's align=True, which aligns them to their component size instead of libcu++'s alignas(2 * sizeof(T)). A RawOp casting to a native struct can read or write the wrong bytes with no error.
RawOp.compile() now checks input and output types against the native layout, including nested structs, subarrays, and pointers, and raises a clear error instead of building a broken kernel.
Fixing the layout itself is bigger separate work: NumPy's packer can't use a different alignment per field, and the MLIR/LLVM codegen for Python callable operators would also need to agree. Those operators aren't affected by this bug either way, since numba-cuda-mlir represents complex values consistently on its own.
Testing
New test_raw_op_complex_layout.py (18 tests), plus test_raw_op.py, test_no_numba.py, test_select.py, test_struct_conversions.py, test_nested_struct.py, and test_struct_field_validation.py all pass.
Checklist
cc @NaderAlAwar for review, and @shwina @Jacobfaib @oleksandr-pavlyk for visibility.