@@ -184,48 +184,52 @@ cdef StridedMemoryView view_as_dlpack(obj, stream_ptr, view=None):
184
184
stream = int (stream_ptr) if stream_ptr else None )
185
185
186
186
cdef void * data = NULL
187
+ cdef DLTensor* dl_tensor
188
+ cdef DLManagedTensorVersioned* dlm_tensor_ver
189
+ cdef DLManagedTensor* dlm_tensor
190
+ cdef const char * used_name
187
191
if cpython.PyCapsule_IsValid(
188
192
capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME):
189
193
data = cpython.PyCapsule_GetPointer(
190
194
capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME)
191
195
versioned = True
196
+ dlm_tensor_ver = < DLManagedTensorVersioned* > data
197
+ dl_tensor = & dlm_tensor_ver.dl_tensor
198
+ is_readonly = bool ((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0 )
199
+ used_name = DLPACK_VERSIONED_TENSOR_USED_NAME
192
200
elif cpython.PyCapsule_IsValid(
193
201
capsule, DLPACK_TENSOR_UNUSED_NAME):
194
202
data = cpython.PyCapsule_GetPointer(
195
203
capsule, DLPACK_TENSOR_UNUSED_NAME)
196
204
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
205
dlm_tensor = < DLManagedTensor* > data
209
206
dl_tensor = & dlm_tensor.dl_tensor
210
207
is_readonly = False
208
+ used_name = DLPACK_TENSOR_USED_NAME
209
+ else :
210
+ assert False
211
211
212
212
cdef StridedMemoryView buf = StridedMemoryView() if view is None else view
213
213
buf.ptr = < intptr_t> (dl_tensor.data)
214
- buf.shape = tuple (int (dl_tensor.shape[i]) for i in range (dl_tensor.ndim))
214
+
215
+ # Construct shape and strides tuples using the Python/C API for speed
216
+ buf.shape = cpython.PyTuple_New(dl_tensor.ndim)
217
+ for i in range (dl_tensor.ndim):
218
+ cpython.PyTuple_SET_ITEM(buf.shape, i, cpython.PyLong_FromLong(dl_tensor.shape[i]))
215
219
if dl_tensor.strides:
216
- buf.strides = tuple (
217
- int (dl_tensor.strides[i]) for i in range (dl_tensor.ndim))
220
+ buf.strides = cpython.PyTuple_New(dl_tensor.ndim)
221
+ for i in range (dl_tensor.ndim):
222
+ cpython.PyTuple_SET_ITEM(buf.strides, i, cpython.PyLong_FromLong(dl_tensor.strides[i]))
218
223
else :
219
224
# C-order
220
225
buf.strides = None
226
+
221
227
buf.dtype = dtype_dlpack_to_numpy(& dl_tensor.dtype)
222
228
buf.device_id = device_id
223
229
buf.is_device_accessible = is_device_accessible
224
230
buf.readonly = is_readonly
225
231
buf.exporting_obj = obj
226
232
227
- cdef const char * used_name = (
228
- DLPACK_VERSIONED_TENSOR_USED_NAME if versioned else DLPACK_TENSOR_USED_NAME)
229
233
cpython.PyCapsule_SetName(capsule, used_name)
230
234
231
235
return buf
@@ -308,8 +312,10 @@ cdef StridedMemoryView view_as_cai(obj, stream_ptr, view=None):
308
312
buf.dtype = numpy.dtype(cai_data[" typestr" ])
309
313
buf.strides = cai_data.get(" strides" )
310
314
if buf.strides is not None :
311
- # convert to counts
312
- buf.strides = tuple (s // buf.dtype.itemsize for s in buf.strides)
315
+ # convert to counts, using the Python/C API for speed
316
+ for i in range (len (buf.strides)):
317
+ cpython.PyTuple_SET_ITEM(buf.strides, i,
318
+ cpython.PyLong_FromLong(buf.strides[i] // buf.dtype.itemsize))
313
319
buf.is_device_accessible = True
314
320
buf.device_id = handle_return(
315
321
driver.cuPointerGetAttribute(
0 commit comments