[cuda.compute] Provide MVP expose of mdspan-based cub::DeviceCopy::Copy
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
# Problem statement
In an ideal world `cuda.compute` would support every Python type used to represent a tensor accessible from CUDA devices, although some types may offer lower host-side overhead.
Some Python types represent fixed-rank tensors, such as `warp.array2d`, `warp.array3d`, etc., and some represent arrays of arbitrary rank specified at run-time, such as `cupy.ndarray`, `tensor.Torch`.
`cuda.compute` needs to be able to convert an instance of supported Python type to an instance of native structure accepted by a host callable that converts its arguments to instances of `mdspan` and calls `cub::DeviceCopy::Copy` ([ref](https://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceCopy.html)).
Under the hood `cuda.compute` may be doing DLPack-based conversion similar in spirit to what [gh:apache/tvm-ffi](https://github.com/apache/tvm-ffi) does to apply user-supplied library function to Python tensor types.
## mdspan C++ type
`mdspan` is a templated type that represents a multi-dimensional array view added in C++23 standard and implemented in CUDA standard C++ library. To instantiate the type one needs to specify element type, array rank as well as array layout at compile time.
1. type communicates size and alignment for elements of type-uniform array
2. rank communicates number of indices needed to uniquely reference logical array element
3. layout translates multi-index logically addressing an array element to its offset within
underlying storage.
C++ standard provides several layouts:
- `layout_left`: representing Fortant-centric layout of arrays in memory, i.e., we start with multi-index origin, vary first index, then second, and so on as we lay their values in underlying storage.
For matrices, we write out column after column so that elements of the same column are adjacent in memory.
- `layout_right`: represent C-centric layout, i.e., we start with multi-index origin, vary last index, then second to last, and so on.
For matrices, we write out row after row so that elements of the same row are adjacent in memory.
- `layout_stride`: represents strided layout that can represent result of slicing (taking sub-sequence of indices) and permutations of an array with contiguous storage. Given multi-index `(i0, i1, ..., i{r-1})` and strides `(s0, s1, ..., s{r-1})`, `storage_offset` is `i0 * s0 + i1 * s1 + ... + i{r-1} * s{r-1}` elements.
\
**Note**: This layout type enjoys nice closure property that slices of an array with this layout can also be represented using `layout_stride` and that axes permutation of an array with this layout can also be represented using `layout_stride`. This imposes a _limitation_ that stride vector must have positive elements, which implies that `storage_offset` is always non-negative.
\
This layout subsumes both `layout_left` and `layout_right` as special cases. For extents array with extends (`N0`, `N1`, ..., `N{r-1}`), the corresponding strides (specified in elements) are:
- `layout_left`: (`1`, `N0`, `N0 * N1`, ..., `N0 * N1 * ... * N{r-2}`)
- `layout_right`: (`N1 * N2 * ... * N{r-1}`, `...`, `N{r-2} * N{r-1}`, `N{r-1}`, `1`)
- `layout_stride_relaxed`: represents layout where strides are allowed to be zero or negative. This layout permits representing arrays where some axes were flipped, i.e. `a(..., i, ...,) := b(...,n-i,...)`. An array with all-zeros stride represents an array with the same element values stored compactly.
Common array implementations in Python, such as NumPy, CuPy, torch, etc., implement strided arrays and naturally work with instances that have negative or zero strides.
The `cub::DeviceCopy::Copy` method expects array views to be specified using `cuda::std::mdpan` templated type.
## DLPack
DLPack is common in-memory tensor structure specification ([gh:dmlc/dlpack](https://github.com/dmlc/dlpack)) inspired by Python's [buffer structure](https://docs.python.org/3/c-api/buffer.html#buffer-structure) but aiming to represent data on accelerator devices which may not always be directly accessible from host.
Very much like `mdspan`, DLPack represents an array view, but rank information is provided as run-time.
DLPack represents rank-r array with strided layout, with strides specified in elements (like in `mdspan`).
DLPack struct encodes array using the following information stored in `DLTensor` struct that encodes:
- void pointer `ptr`
- byte-offset from `ptr` to the start of data
- rank value (non-negative)
- pointer to int64 array of `rank` positive extents
- pointer to int64 array of `rank` strides (if nullptr, array is assumed C-contiguous)
- device where data resides, encoded as enum describing vendor + allocation-type (kdCUDA, kdCUDAPinned, etc.), and device ordinal.
Use of byte-offset permits describing an array backed by an unaligned allocation, with (a non-negative) offset used to achieved alignment requirements for `ptr + byte_offset` address of the array value with zero multi-index.
Byte-offset may also be used to guarantee that even with negative strides the total offset `byte_offset + (i0 * s0 + i1 * s1 + ...) * itemsize` is not negative for valid multi-indices.
DLPack struct may encode `ptr` as the address of original device allocation and use `byte_offset` to indicate where
array data start, but this is not specified, and implementors often use `byte_offset = 0` and manipulate `ptr` instead.
DLPack is officially supported by Python Array API standard, see https://data-apis.org/array-api/latest/.
CUB provides a utility function to validate `DLTensor` and convert it to a fully qualified `mdspan` type.
## Performance considerations
We should be conscious of the host-side overhead of converting Python type to the underlying array view structure. The best possible performance is achieved for a hypothetical Python array type that stores array view within the Python object itself so that the conversion is simply a pointer lookup.
Generically, like TVM-FFI does, conversion would go through DLPack. `cuda.compute` checks whether objects supplied as tensor arguments support `__dlpack_c_exchange_api__` or `__dlpack__` protocol. If they do, DLTensor is rearranged to create the view.
Use of `__dlpack_c_exchange_api__` dramatically reduces the overhead, although it still remains significant relative to the cost of conversion of the hypothetical Python array type mentioned above.
Contributor guide
Assessment
This issue has not been assessed yet.