@@ -12,6 +12,7 @@ from typing import Any, Optional
12
12
import numpy
13
13
14
14
from cuda.core.experimental._utils.cuda_utils import handle_return, driver
15
+ from cuda.core.experimental._utils cimport cuda_utils
15
16
16
17
17
18
# TODO(leofang): support NumPy structured dtypes
@@ -184,48 +185,47 @@ cdef StridedMemoryView view_as_dlpack(obj, stream_ptr, view=None):
184
185
stream = int (stream_ptr) if stream_ptr else None )
185
186
186
187
cdef void * data = NULL
188
+ cdef DLTensor* dl_tensor
189
+ cdef DLManagedTensorVersioned* dlm_tensor_ver
190
+ cdef DLManagedTensor* dlm_tensor
191
+ cdef const char * used_name
187
192
if cpython.PyCapsule_IsValid(
188
193
capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME):
189
194
data = cpython.PyCapsule_GetPointer(
190
195
capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME)
191
196
versioned = True
197
+ dlm_tensor_ver = < DLManagedTensorVersioned* > data
198
+ dl_tensor = & dlm_tensor_ver.dl_tensor
199
+ is_readonly = bool ((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0 )
200
+ used_name = DLPACK_VERSIONED_TENSOR_USED_NAME
192
201
elif cpython.PyCapsule_IsValid(
193
202
capsule, DLPACK_TENSOR_UNUSED_NAME):
194
203
data = cpython.PyCapsule_GetPointer(
195
204
capsule, DLPACK_TENSOR_UNUSED_NAME)
196
205
versioned = False
197
- else :
198
- assert False
199
-
200
- cdef DLManagedTensor* dlm_tensor
201
- cdef DLManagedTensorVersioned* dlm_tensor_ver
202
- cdef DLTensor* dl_tensor
203
- if versioned:
204
- dlm_tensor_ver = < DLManagedTensorVersioned* > data
205
- dl_tensor = & dlm_tensor_ver.dl_tensor
206
- is_readonly = bool ((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0 )
207
- else :
208
206
dlm_tensor = < DLManagedTensor* > data
209
207
dl_tensor = & dlm_tensor.dl_tensor
210
208
is_readonly = False
209
+ used_name = DLPACK_TENSOR_USED_NAME
210
+ else :
211
+ assert False
211
212
212
213
cdef StridedMemoryView buf = StridedMemoryView() if view is None else view
213
214
buf.ptr = < intptr_t> (dl_tensor.data)
214
- buf.shape = tuple (int (dl_tensor.shape[i]) for i in range (dl_tensor.ndim))
215
+
216
+ buf.shape = cuda_utils.carray_int64_t_to_tuple(dl_tensor.shape, dl_tensor.ndim)
215
217
if dl_tensor.strides:
216
- buf.strides = tuple (
217
- int (dl_tensor.strides[i]) for i in range (dl_tensor.ndim))
218
+ buf.strides = cuda_utils.carray_int64_t_to_tuple(dl_tensor.strides, dl_tensor.ndim)
218
219
else :
219
220
# C-order
220
221
buf.strides = None
222
+
221
223
buf.dtype = dtype_dlpack_to_numpy(& dl_tensor.dtype)
222
224
buf.device_id = device_id
223
225
buf.is_device_accessible = is_device_accessible
224
226
buf.readonly = is_readonly
225
227
buf.exporting_obj = obj
226
228
227
- cdef const char * used_name = (
228
- DLPACK_VERSIONED_TENSOR_USED_NAME if versioned else DLPACK_TENSOR_USED_NAME)
229
229
cpython.PyCapsule_SetName(capsule, used_name)
230
230
231
231
return buf
0 commit comments