feat: relax the power of two check in StridedLayout (NVIDIA#1427) (NVIDIA#1471)

Co-authored-by: Phillip Cloud <417981+cpcloud@users.noreply.github.com>

Author: github-actions[bot]

cuda_core/cuda/core/_layout.pxd

@@ -111,7 +111,8 @@ cdef class _StridedLayout:
     # ==============================
 
     cdef inline int _init(_StridedLayout self, BaseLayout& base, int itemsize, bint divide_strides=False) except -1 nogil:
-        _validate_itemsize(itemsize)
+        if itemsize <= 0:
+            raise ValueError("itemsize must be positive")
 
         if base.strides != NULL and divide_strides:
             _divide_strides(base, itemsize)
@@ -123,7 +124,8 @@ cdef class _StridedLayout:
         return 0
 
     cdef inline stride_t _init_dense(_StridedLayout self, BaseLayout& base, int itemsize, OrderFlag order_flag, axis_vec_t* stride_order=NULL) except -1 nogil:
-        _validate_itemsize(itemsize)
+        if itemsize <= 0:
+            raise ValueError("itemsize must be positive")
 
         cdef stride_t volume
         if order_flag == ORDER_C:
@@ -643,14 +645,6 @@ cdef inline bint _normalize_axis(integer_t& axis, integer_t extent) except -1 no
     return True
 
 
-cdef inline int _validate_itemsize(int itemsize) except -1 nogil:
-    if itemsize <= 0:
-        raise ValueError("itemsize must be positive")
-    if itemsize & (itemsize - 1):
-        raise ValueError("itemsize must be a power of two")
-    return 0
-
-
 cdef inline bint _is_unique(BaseLayout& base, axis_vec_t& stride_order) except -1 nogil:
     if base.strides == NULL:
         return True

cuda_core/cuda/core/_layout.pyx

@@ -29,7 +29,7 @@ cdef class _StridedLayout:
         Otherwise, the strides are assumed to be implicitly C-contiguous and the resulting
         layout's :attr:`strides` will be None.
     itemsize : int
-        The number of bytes per single element (dtype size). Must be a power of two.
+        The number of bytes per single element (dtype size).
     divide_strides : bool, optional
         If True, the provided :attr:`strides` will be divided by the :attr:`itemsize`.
 
@@ -40,7 +40,7 @@ cdef class _StridedLayout:
     Attributes
     ----------
     itemsize : int
-        The number of bytes per single element (dtype size). Must be a power of two.
+        The number of bytes per single element (dtype size).
     slice_offset : int
         The offset (as a number of elements, not bytes) of the element at
         index ``(0,) * ndim``. See also :attr:`slice_offset_in_bytes`.
@@ -636,7 +636,6 @@ cdef class _StridedLayout:
         In either case, the ``volume * itemsize`` of the layout remains the same.
 
         The conversion is subject to the following constraints:
-            * The old and new itemsizes must be powers of two.
             * The extent at ``axis`` must be a positive integer.
             * The stride at ``axis`` must be 1.
 
@@ -1214,10 +1213,10 @@ cdef inline int64_t gcd(int64_t a, int64_t b) except? -1 nogil:
 
 cdef inline int pack_extents(BaseLayout& out_layout, stride_t& out_slice_offset, BaseLayout& in_layout, stride_t slice_offset, int itemsize, int new_itemsize, intptr_t data_ptr, bint keep_dim, int axis) except -1 nogil:
     cdef int ndim = in_layout.ndim
-    if new_itemsize <= 0 or new_itemsize & (new_itemsize - 1):
-        raise ValueError(f"new itemsize must be a power of two, got {new_itemsize}.")
-    if itemsize <= 0 or itemsize & (itemsize - 1):
-        raise ValueError(f"itemsize must be a power of two, got {itemsize}.")
+    if new_itemsize <= 0:
+        raise ValueError(f"new itemsize must be greater than zero, got {new_itemsize}.")
+    if itemsize <= 0:
+        raise ValueError(f"itemsize must be greater than zero, got {itemsize}.")
     if new_itemsize <= itemsize:
         if new_itemsize == itemsize:
             return 1
@@ -1270,10 +1269,10 @@ cdef inline int unpack_extents(BaseLayout &out_layout, BaseLayout &in_layout, in
     cdef int ndim = in_layout.ndim
     if not _normalize_axis(axis, ndim):
         raise ValueError(f"Invalid axis: {axis} out of range for {ndim}D tensor")
-    if new_itemsize <= 0 or new_itemsize & (new_itemsize - 1):
-        raise ValueError(f"new itemsize must be a power of two, got {new_itemsize}.")
-    if itemsize <= 0 or itemsize & (itemsize - 1):
-        raise ValueError(f"itemsize must be a power of two, got {itemsize}.")
+    if new_itemsize <= 0:
+        raise ValueError(f"new itemsize must be greater than zero, got {new_itemsize}.")
+    if itemsize <= 0:
+        raise ValueError(f"itemsize must be greater than zero, got {itemsize}.")
     if new_itemsize >= itemsize:
         if new_itemsize == itemsize:
             return 1
@@ -1301,10 +1300,10 @@ cdef inline int unpack_extents(BaseLayout &out_layout, BaseLayout &in_layout, in
 
 cdef inline int max_compatible_itemsize(BaseLayout& layout, stride_t slice_offset, int itemsize, int max_itemsize, intptr_t data_ptr, int axis) except? -1 nogil:
     cdef int ndim = layout.ndim
-    if max_itemsize <= 0 or max_itemsize & (max_itemsize - 1):
-        raise ValueError(f"max_itemsize must be a power of two, got {max_itemsize}.")
-    if itemsize <= 0 or itemsize & (itemsize - 1):
-        raise ValueError(f"itemsize must be a power of two, got {itemsize}.")
+    if max_itemsize <= 0:
+        raise ValueError(f"max_itemsize must be greater than zero, got {max_itemsize}.")
+    if itemsize <= 0:
+        raise ValueError(f"itemsize must be greater than zero, got {itemsize}.")
     if not _normalize_axis(axis, ndim):
         raise ValueError(f"Invalid axis: {axis} out of range for {ndim}D tensor")
     if max_itemsize < itemsize:

cuda_core/tests/test_utils.py

@@ -433,3 +433,16 @@ def test_view_zero_size_array(api, shape, dtype):
     assert smv.size == 0
     assert smv.shape == shape
     assert smv.dtype == np.dtype(dtype)
+
+
+def test_from_buffer_with_non_power_of_two_itemsize():
+    dev = Device()
+    dev.set_current()
+    dtype = np.dtype([("a", "int32"), ("b", "int8")])
+    shape = (1,)
+    layout = _StridedLayout(shape=shape, strides=None, itemsize=dtype.itemsize)
+    required_size = layout.required_size_in_bytes()
+    assert required_size == math.prod(shape) * dtype.itemsize
+    buffer = dev.memory_resource.allocate(required_size)
+    view = StridedMemoryView.from_buffer(buffer, shape=shape, strides=layout.strides, dtype=dtype, is_readonly=True)
+    assert view.dtype == dtype