Skip to content

Commit

Permalink
Versioned DLpack tensor deleter and owner tweaked
Browse files Browse the repository at this point in the history
These tweaks are based on similar changes made for the managed tensor during memory work
  • Loading branch information
ndgrigorian committed May 10, 2024
1 parent a7d255a commit c4b34e5
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions dpctl/tensor/_dlpack.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -160,8 +160,10 @@ cdef void _pycapsule_versioned_deleter(object dlt_capsule) noexcept:

cdef void _managed_tensor_versioned_deleter(DLManagedTensorVersioned *dlmv_tensor) noexcept with gil:
if dlmv_tensor is not NULL:
# we only delete shape, because we make single allocation to
# acommodate both shape and strides if strides are needed
stdlib.free(dlmv_tensor.dl_tensor.shape)
cpython.Py_DECREF(<usm_ndarray>dlmv_tensor.manager_ctx)
cpython.Py_DECREF(<object>dlmv_tensor.manager_ctx)
dlmv_tensor.manager_ctx = NULL
stdlib.free(dlmv_tensor)

Expand Down Expand Up @@ -543,14 +545,15 @@ cdef class _DLManagedTensorVersionedOwner:
Helper class managing the lifetime of the DLManagedTensorVersioned
struct transferred from a 'dlpack_versioned' capsule.
"""
cdef DLManagedTensorVersioned *dlmv_tensor
cdef DLManagedTensorVersioned * dlmv_tensor

def __cinit__(self):
self.dlmv_tensor = NULL

def __dealloc__(self):
if self.dlmv_tensor:
self.dlmv_tensor.deleter(self.dlmv_tensor)
self.dlmv_tensor = NULL

@staticmethod
cdef _DLManagedTensorVersionedOwner _create(DLManagedTensorVersioned *dlmv_tensor_src):
Expand Down

0 comments on commit c4b34e5

Please sign in to comment.