Skip to content

Prepare mdspan for layout_stride_relaxed support - #9877

Draft
leofang wants to merge 7 commits into
cupy:mainfrom
leofang:layout_stride_relaxed_prep
Draft

leofang wants to merge 7 commits into
cupy:mainfrom
leofang:layout_stride_relaxed_prep

Conversation

@leofang

@leofang leofang commented Apr 18, 2026 •

Copy link
Copy Markdown
Member

Summary

Add support for CCCL's cuda::layout_stride_relaxed to CuPy's mdspan, enabling correct handling of negative strides (reversed array views) and zero strides (broadcasting).

Struct layout

The packed mdspan struct now includes an offset field:

[void* ptr | IndexType extents[N] | IndexType strides[N] | ptrdiff_t offset]

The offset is always ptrdiff_t (8 bytes) to match layout_stride_relaxed::mapping's default offset_type. This is important for the IndexType=int32 case where the extents and strides are 4 bytes each but the offset remains 8 bytes:

Component IndexType=int32 IndexType=int64
ptr 8 B 8 B
extents[N] N × 4 B N × 8 B
strides[N] N × 4 B N × 8 B
offset 8 B (ptrdiff_t) 8 B (ptrdiff_t)

xref: https://github.com/NVIDIA/cccl/blob/dee407140bde57111dd1fa9311434cb26429f035/libcudacxx/include/cuda/__fwd/mdspan.h#L82-L85

Pointer and offset semantics

  • Non-negative strides (or allow_unsafe=True): ptr = logical data pointer, offset = 0. Backward compatible with layout_stride, layout_left, layout_right (they ignore the trailing offset).
  • Negative strides with allow_unsafe=False (default): ptr = base of allocation (arr.data.mem.ptr), offset = element distance from base to logical origin. The kernel must use cuda::layout_stride_relaxed.
  • allow_unsafe=True with negative strides: legacy convention (ptr=data.ptr, offset=0) for backward compatibility with layout_stride kernels.

Negative strides no longer require allow_unsafe=True. Only zero strides and zero-size dimensions still do.

Other changes

  • Remove NumericTraits specializations for complex types — they declared complex<float/double> as primitive (is_primitive=true), causing UB in CUB's decoupled lookback scan (torn reads/writes). CUB reductions still work correctly without these traits. (Fixes [BUG] ScanTileState static_assert fails for types with is_primitive=true and sizeof >= 16 NVIDIA/cccl#8207)
  • Bump CCCL submodule to dee407140b (includes layout_stride_relaxed)
  • Add TestMdspanStrideRelaxed test class with negative-stride kernel tests
  • Update validation test to use layout_stride_relaxed with correct sizeof check

Merge blocker

The CCCL submodule currently points to a main commit (dee407140b), not a stable release. Before merging, it should be updated to the first stable CCCL release that includes layout_stride_relaxed. Bumping CCCL is disruptive (invalidates CI cache) so we want to do it only once.

Test plan

  • All 108 mdspan tests pass (0 failures, 0 skips)
  • CUB reduction tests pass for complex types after NumericTraits removal (80/80)
  • Existing mdspan tests (0D, 1D, 2D, 3D with all layouts and dtypes)
  • New TestMdspanStrideRelaxed tests (1D/2D negative strides, both-axes reversed, sliced+reversed)
  • Validation: negative strides accepted by default, zero strides rejected
  • Full CI pass

Closes #9772

-- Leo's bot

leofang and others added 7 commits March 28, 2026 18:35
Add an offset field to the mdspan struct layout to support CCCL's
upcoming cuda::layout_stride_relaxed, which handles negative strides
correctly. The new struct layout is:
  [ptr | extents[N] | strides[N] | offset]

For arrays with non-negative strides (or allow_unsafe=True), ptr is
the logical data pointer and offset is 0, maintaining backward
compatibility with layout_stride/layout_left/layout_right.

For arrays with negative strides and allow_unsafe=False (default),
ptr is the base of the allocation and offset is computed so that the
layout mapping always returns a non-negative index. This makes
negative strides safe by default when used with layout_stride_relaxed
on the device side.

Negative strides no longer require allow_unsafe=True. Only zero
strides and zero-size dimensions still require it.

See: cupy#9772

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
The NumericTraits specializations for complex<float> and complex<double>
marked them as primitive (is_primitive=true), which caused UB in CUB's
decoupled lookback scan due to torn reads/writes when sizeof(T) >= 8.
The CUB scan path for complex128 was already disabled at the Python
level (cub.pyx:514) since 2020 as a workaround for this very issue.

CUB device reductions (sum, min, max, argmin, argmax) and scans
continue to work correctly for complex types without these traits,
as CuPy provides custom operator specializations and
cuda::std::numeric_limits/is_floating_point_v specializations.

This also unblocks upgrading the CCCL submodule to versions containing
layout_stride_relaxed, where the static_assert in ScanTileState
(introduced by NVIDIA/cccl#7575) would fail for is_primitive types
with sizeof >= 16.

Closes: NVIDIA/cccl#8207

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Update the CCCL submodule to dee407140b on main, which includes
cuda::layout_stride_relaxed (NVIDIA/cccl#7175). This layout policy
supports negative and zero strides with an explicit offset, enabling
correct mdspan interop for reversed/strided array views.

NOTE: This commit points to a CCCL main commit, not a stable release.
Before merging this PR, the submodule should be updated to the first
stable CCCL release that includes layout_stride_relaxed. Bumping CCCL
is disruptive (invalidates CI cache), so we want to do it only once.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
The offset field must be packed as ptrdiff_t (8 bytes), not IndexType,
to match the C++ layout_stride_relaxed::mapping's default offset_type.
This matters for IndexType=int32 where sizeof(int)=4 but the C++ struct
expects an 8-byte offset field.

Struct layout (2D example, IndexType=int32):
  [void* ptr (8B)] [int ext0, ext1 (8B)] [int str0, str1 (8B)] [ptrdiff_t offset (8B)]
  Total: 32 bytes = sizeof(void*) + 4*sizeof(int) + sizeof(ptrdiff_t)

Struct layout (2D example, IndexType=int64):
  [void* ptr (8B)] [int64 ext0, ext1 (16B)] [int64 str0, str1 (16B)] [ptrdiff_t offset (8B)]
  Total: 48 bytes = sizeof(void*) + 4*sizeof(int64) + sizeof(ptrdiff_t)

Also fixes:
- _has_layout_stride_relaxed() path to use the installed CCCL symlinks
- Validation test to use layout_stride_relaxed with correct sizeof
- Add skip_no_relaxed_layout to TestMdspanValidation

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
- cimport ptrdiff_t from libc.stddef and use it instead of long long
  for the offset field in _carray.pyx
- Restore original negative-stride tests with layout_stride +
  allow_unsafe=True (legacy path)
- Add separate _relaxed test variants with layout_stride_relaxed
  (safe path, skip_no_relaxed_layout)
- Split TestMdspanValidation into layout_stride (always runs) and
  TestMdspanValidationRelaxed (skip_no_relaxed_layout)
- skip_no_relaxed_layout only applied to tests that exclusively
  target layout_stride_relaxed

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
- Add layout_stride_relaxed to the layout fixture via pytest.param
  with skip_no_relaxed_layout mark, so all existing tests automatically
  cover the relaxed layout
- Merge negative-stride tests: use allow_unsafe conditionally based on
  layout (layout_stride=True, layout_stride_relaxed=False)
- Add both-axes-reversed and sliced-negative-stride tests to
  TestMdspan2D (using the fixture)
- Remove redundant TestMdspanStrideRelaxed class and
  _make_relaxed_kernel_module helper
- Update skip conditions to allow both strided layout variants

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
@leofang leofang added the st:blocked-by-another-pr Blocked by another pull-request label Apr 20, 2026
@leofang

leofang commented Apr 20, 2026

Copy link
Copy Markdown
Member Author

This PR is blocked until there is a public CCCL release that includes layout_stride_relaxed.

Comment thread cupy/_core/_carray.pyx
# returns a non-negative index.
stored_ptr = mem_ptr
mem_offset = (
(<intptr_t>data_ptr - <intptr_t>mem_ptr) // itemsize

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we trust that mem_ptr is a base ptr? As an example, packages don't bother to distinguish between data ptr and base ptr when exporting memory with dlpack, so if the array converted here comes from elsewhere - this won't give us the offset. The elsewhere can by cupy itself:

>>> x = cp.arange(3 * 5, dtype=cp.int8)[::-1]
>>> x.data.ptr - x.data.mem.ptr
14
>>> y = cp.from_dlpack(x.__dlpack__())
>>> y.data.ptr - y.data.mem.ptr  # both point at the last element of the array
0

Idea: ignore data.mem.ptr. Instead, compute min_offset=sum((e - 1) * s for e, s in zip(shape, strides) if s < 0 and e > 0) (note it is zero or negative), then take offset=-min_offset and mem_ptr = data.ptr + min_offset.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

st:blocked-by-another-pr Blocked by another pull-request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Prepare to support mdspan's layout_stride_relaxed [BUG] ScanTileState static_assert fails for types with is_primitive=true and sizeof >= 16

2 participants