Revert dlpack changes

Signed-off-by: Kamil Tokarski <ktokarski@nvidia.com>

Author: stiepan

cuda_core/cuda/core/experimental/_dlpack.pyx

@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import IntEnum
-from ._layout cimport StridedLayout
 
 
 cdef void pycapsule_deleter(object capsule) noexcept:
@@ -34,16 +33,6 @@ cdef void deleter(DLManagedTensor* tensor) noexcept with gil:
     stdlib.free(tensor)
 
 
-cdef void cleanup(DLManagedTensor* tensor) noexcept with gil:
-    if tensor:
-        if tensor.dl_tensor.shape:
-            stdlib.free(tensor.dl_tensor.shape)
-        if tensor.manager_ctx:
-            cpython.Py_DECREF(<object>tensor.manager_ctx)
-            tensor.manager_ctx = NULL
-        stdlib.free(tensor)
-
-
 cdef void versioned_deleter(DLManagedTensorVersioned* tensor) noexcept with gil:
     stdlib.free(tensor.dl_tensor.shape)
     if tensor.manager_ctx:
@@ -52,17 +41,43 @@ cdef void versioned_deleter(DLManagedTensorVersioned* tensor) noexcept with gil:
     stdlib.free(tensor)
 
 
-cdef void cleanup_versioned(DLManagedTensorVersioned* tensor) noexcept with gil:
-    if tensor:
-        if tensor.dl_tensor.shape:
-            stdlib.free(tensor.dl_tensor.shape)
-        if tensor.manager_ctx:
-            cpython.Py_DECREF(<object>tensor.manager_ctx)
-            tensor.manager_ctx = NULL
-        stdlib.free(tensor)
+cpdef object make_py_capsule(object buf, bint versioned):
+    cdef DLManagedTensor* dlm_tensor
+    cdef DLManagedTensorVersioned* dlm_tensor_ver
+    cdef DLTensor* dl_tensor
+    cdef void* tensor_ptr
+    cdef const char* capsule_name
 
+    if versioned:
+        dlm_tensor_ver = <DLManagedTensorVersioned*>(
+            stdlib.malloc(sizeof(DLManagedTensorVersioned)))
+        dlm_tensor_ver.version.major = DLPACK_MAJOR_VERSION
+        dlm_tensor_ver.version.minor = DLPACK_MINOR_VERSION
+        dlm_tensor_ver.manager_ctx = <void*>buf
+        dlm_tensor_ver.deleter = versioned_deleter
+        dlm_tensor_ver.flags = 0
+        dl_tensor = &dlm_tensor_ver.dl_tensor
+        tensor_ptr = dlm_tensor_ver
+        capsule_name = DLPACK_VERSIONED_TENSOR_UNUSED_NAME
+    else:
+        dlm_tensor = <DLManagedTensor*>(
+            stdlib.malloc(sizeof(DLManagedTensor)))
+        dl_tensor = &dlm_tensor.dl_tensor
+        dlm_tensor.manager_ctx = <void*>buf
+        dlm_tensor.deleter = deleter
+        tensor_ptr = dlm_tensor
+        capsule_name = DLPACK_TENSOR_UNUSED_NAME
+
+    dl_tensor.data = <void*><intptr_t>(int(buf.handle))
+    dl_tensor.ndim = 1
+    cdef int64_t* shape_strides = \
+        <int64_t*>stdlib.malloc(sizeof(int64_t) * 2)
+    shape_strides[0] = <int64_t>buf.size
+    shape_strides[1] = 1  # redundant
+    dl_tensor.shape = shape_strides
+    dl_tensor.strides = NULL
+    dl_tensor.byte_offset = 0
 
-cdef inline int _setup_dl_tensor_device(DLTensor* dl_tensor, object buf) except -1:
     cdef DLDevice* device = &dl_tensor.device
     # buf should be a Buffer instance
     if buf.is_device_accessible and not buf.is_host_accessible:
@@ -76,103 +91,14 @@ cdef inline int _setup_dl_tensor_device(DLTensor* dl_tensor, object buf) except
         device.device_id = 0
     else:  # not buf.is_device_accessible and not buf.is_host_accessible
         raise BufferError("invalid buffer")
-    return 0
-
-
-cdef inline int _setup_dl_tensor_layout(DLTensor* dl_tensor, object buf, StridedLayout layout) except -1:
-    cdef int64_t* shape_strides = NULL
-    cdef int ndim
-    if layout is None:
-        dl_tensor.ndim = 1
-        shape_strides = <int64_t*>stdlib.malloc(sizeof(int64_t) * 2)
-        shape_strides[0] = <int64_t>buf.size
-        shape_strides[1] = 1  # redundant
-        dl_tensor.shape = shape_strides
-        dl_tensor.strides = NULL
-        dl_tensor.byte_offset = 0
-    else:
-        ndim = layout.ndim
-        dl_tensor.ndim = ndim
-        shape_strides = <int64_t*>stdlib.malloc(sizeof(int64_t) * ndim * 2)
-        dl_tensor.shape = shape_strides
-        for i in range(ndim):
-            shape_strides[i] = layout.base.shape[i]
-        if layout.base.strides == NULL:
-            dl_tensor.strides = NULL
-        else:
-            dl_tensor.strides = shape_strides + ndim
-            for i in range(ndim):
-                dl_tensor.strides[i] = layout.base.strides[i]
-        dl_tensor.byte_offset = 0
-    return 0
-
-
-cdef inline int _setup_dl_tensor_dtype(DLTensor* dl_tensor, object dtype) except -1:
-    cdef DLDataType* dl_dtype = &dl_tensor.dtype
-    if dtype is None:
-        dl_dtype.code = <uint8_t>kDLInt
-        dl_dtype.lanes = <uint16_t>1
-        dl_dtype.bits = <uint8_t>8
-        return 0
-    cdef uint8_t code
-    cdef uint8_t bits
-    cdef uint16_t lanes
-    code, bits, lanes = dtype
-    dl_dtype.code = code
-    dl_dtype.bits = bits
-    dl_dtype.lanes = lanes
-    return 0
-
-
-cpdef object make_py_capsule(
-    object buf,
-    bint versioned,
-    intptr_t data_ptr,
-    StridedLayout layout=None,
-    object dtype=None,
-):
-    cdef DLManagedTensor* dlm_tensor = NULL
-    cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL
-    cdef DLTensor* dl_tensor = NULL
-    cdef void* tensor_ptr
-    cdef const char* capsule_name
 
-    try:
-        if versioned:
-            dlm_tensor_ver = <DLManagedTensorVersioned*>(
-                stdlib.malloc(sizeof(DLManagedTensorVersioned)))
-            dlm_tensor_ver.dl_tensor.shape = NULL
-            dl_tensor = &dlm_tensor_ver.dl_tensor
-            dlm_tensor_ver.version.major = DLPACK_MAJOR_VERSION
-            dlm_tensor_ver.version.minor = DLPACK_MINOR_VERSION
-            cpython.Py_INCREF(buf)
-            dlm_tensor_ver.manager_ctx = <void*>buf
-            dlm_tensor_ver.deleter = versioned_deleter
-            dlm_tensor_ver.flags = 0
-            tensor_ptr = dlm_tensor_ver
-            capsule_name = DLPACK_VERSIONED_TENSOR_UNUSED_NAME
-        else:
-            dlm_tensor = <DLManagedTensor*>(
-                stdlib.malloc(sizeof(DLManagedTensor)))
-            dlm_tensor.dl_tensor.shape = NULL
-            dl_tensor = &dlm_tensor.dl_tensor
-            cpython.Py_INCREF(buf)
-            dlm_tensor.manager_ctx = <void*>buf
-            dlm_tensor.deleter = deleter
-            tensor_ptr = dlm_tensor
-            capsule_name = DLPACK_TENSOR_UNUSED_NAME
-
-        dl_tensor.data = <void*>data_ptr
-
-        _setup_dl_tensor_device(dl_tensor, buf)
-        _setup_dl_tensor_layout(dl_tensor, buf, layout)
-        _setup_dl_tensor_dtype(dl_tensor, dtype)
-
-        return cpython.PyCapsule_New(tensor_ptr, capsule_name, pycapsule_deleter)
-    except:
-        cleanup(dlm_tensor)
-        cleanup_versioned(dlm_tensor_ver)
-        raise
+    cdef DLDataType* dtype = &dl_tensor.dtype
+    dtype.code = <uint8_t>kDLInt
+    dtype.lanes = <uint16_t>1
+    dtype.bits = <uint8_t>8
+
+    cpython.Py_INCREF(buf)
+    return cpython.PyCapsule_New(tensor_ptr, capsule_name, pycapsule_deleter)
 
 
 class DLDeviceType(IntEnum):

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

@@ -203,7 +203,7 @@ cdef class Buffer:
             if not isinstance(max_version, tuple) or len(max_version) != 2:
                 raise BufferError(f"Expected max_version tuple[int, int], got {max_version}")
             versioned = max_version >= (1, 0)
-        capsule = make_py_capsule(self, versioned, int(self.handle))
+        capsule = make_py_capsule(self, versioned)
         return capsule
 
     def __dlpack_device__(self) -> tuple[int, int]: