Coverage for cuda/core/_memoryview.pyx: 85.07%

1# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. (empty)

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

7from ._dlpack cimport * 

8from ._dlpack import classify_dl_device (empty)

9from libc.stdint cimport intptr_t 

10from cuda.core._layout cimport _StridedLayout, get_strides_ptr 

11from cuda.core._stream import Stream (empty)

12  

13import ctypes (empty)

14import functools (empty)

15import sys (empty)

16import warnings (empty)

17from collections.abc import Callable # no-cython-lint # used in string annotations below (empty)

18from typing import Any # no-cython-lint # used in string annotations below (empty)

19  

20import numpy (empty)

21  

22from cuda.bindings cimport cydriver 

23from cuda.core._resource_handles cimport ( 

24 EventHandle, 

25 create_event_handle_noctx, 

26 as_cu, 

27) 

28  

29from cuda.core._utils.cuda_utils import handle_return, driver (empty)

30from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

31  

32  

33from cuda.core._memory import Buffer (empty)

34  

35  

36# --------------------------------------------------------------------------- 

37# Lazy tensor bridge (avoids loading _tensor_bridge.so until torch is used) 

38# --------------------------------------------------------------------------- 

39  

40cdef object _tensor_bridge = None (empty)

41# Cache: type(obj) -> True/False for the torch tensor check. 

42# Once a type is seen, we never re-check. 

43cdef dict _torch_type_cache = {} (empty)

44# Tri-state: None = not checked, True/False = result of version check 

45cdef object _torch_version_ok = None (empty)

46  

47cdef inline bint _torch_version_check(): (empty)

48 """Return True if 2.3 <= torch <= 2.12 (known AOTI ABI range). Memoized. 

49  

50 Lower bound: AOTI functions we use were introduced in PyTorch 2.3. 

51 Upper bound: the ``pyobj_to_aten_handle`` trick relies on the 

52 THPVariable struct layout (PyObject_HEAD followed by at::Tensor cdata) 

53 and the identity ``AtenTensorHandle == at::Tensor*``. Both are 

54 undocumented internals that could change in a future PyTorch version. 

55 We cap at the latest version we have tested against; unknown versions 

56 fall back to the standard DLPack/CAI paths. Bump the upper bound 

57 after verifying a new PyTorch release. 

58 """ 

59 global _torch_version_ok 

60 if _torch_version_ok is not None: 

61 return <bint>_torch_version_ok 

62 torch = sys.modules.get("torch") 

63 if torch is None: 

64 _torch_version_ok = False 

65 return False 

66 try: 

67 major, minor = int(torch.__version__.split(".")[0]), \ 

68 int(torch.__version__.split(".")[1]) 

69 _torch_version_ok = (2, 3) <= (major, minor) <= (2, 12) 

70 except (ValueError, IndexError): 

71 _torch_version_ok = False 

72 return <bint>_torch_version_ok 

73  

74  

75cdef inline bint _is_torch_tensor(object obj): (empty)

76 cdef type tp = type(obj) 84 ctx21 2 3 8 V P J K L M N Q R S T U v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb| Ib~ 6 7 c I 5 H 4 B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

77 cdef object cached = _torch_type_cache.get(tp) 84 ctx21 2 3 8 V P J K L M N Q R S T U v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb| Ib~ 6 7 c I 5 H 4 B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

78 if cached is not None: 84 ctx21 2 3 8 V P J K L M N Q R S T U v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb| Ib~ 6 7 c I 5 H 4 B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

79 return <bint>cached 75 ctx23 8 V P J K L M N Q R S T U x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb6 7 I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba

80 cdef str mod = tp.__module__ or "" 10 ctx21 2 v | Ib~ c 4 a Mb

81 cdef bint result = mod.startswith("torch") and hasattr(obj, "data_ptr") \ 10 ctx21 2 v | Ib~ c 4 a Mb

82 and _torch_version_check() 

83 _torch_type_cache[tp] = result # setdefault not needed for bools 10 ctx21 2 v | Ib~ c 4 a Mb

84 return result 10 ctx21 2 v | Ib~ c 4 a Mb

85  

86  

87cdef object _get_tensor_bridge(): (empty)

88 """Bootstrap AOTI symbols, then import _tensor_bridge on first use.""" 

89 global _tensor_bridge 

90 if _tensor_bridge is not None: 

91 return _tensor_bridge 

92 torch_C = sys.modules.get("torch._C") 

93 if torch_C is None: 

94 raise RuntimeError( 

95 "torch._C is not loaded; cannot initialise the tensor bridge. " 

96 "Make sure PyTorch is imported before passing a torch.Tensor.") 

97 ctypes.CDLL(torch_C.__file__, mode=ctypes.RTLD_GLOBAL) 

98 from cuda.core import _tensor_bridge as tb 

99 _tensor_bridge = tb 

100 return _tensor_bridge 

101  

102  

103try: (empty)

104 from ml_dtypes import bfloat16 (empty)

105except ImportError: (empty)

106 bfloat16 = None (empty)

107  

108# TODO(leofang): support NumPy structured dtypes 

109  

110  

111cdef extern from "Python.h": 

112 ctypedef struct PyTypeObject: 

113 void* tp_dict 

114 void PyType_Modified(PyTypeObject*) 

115  

116  

117cdef DLPackExchangeAPI _SMV_DLPACK_EXCHANGE_API 

118cdef bint _SMV_DLPACK_EXCHANGE_API_INITED = False (empty)

119_SMV_DLPACK_EXCHANGE_API_CAPSULE = cpython.PyCapsule_New( (empty)

120 <void*>&_SMV_DLPACK_EXCHANGE_API, 

121 b"dlpack_exchange_api", 

122 NULL, 

123) 

124  

125  

126cdef class StridedMemoryView: 

127 """A class holding metadata of a strided dense array/tensor. 

128  

129 A :obj:`StridedMemoryView` instance can be created in three ways: 

130  

131 1. Using the :obj:`args_viewable_as_strided_memory` decorator (recommended) 

132 2. Explicit construction relying on DLPack or CUDA Array Interface, see below. 

133 3. From :obj:`~_memory.Buffer` and shape and size tuples (see 

134 :meth:`from_buffer` classmethod) 

135  

136 ``StridedMemoryView(obj, stream_ptr)`` can be used to create a view from 

137 objects supporting either DLPack (up to v1.0) or CUDA Array Interface 

138 (CAI) v3. When wrapping an arbitrary object it will try the DLPack protocol 

139 first, then the CAI protocol. A :obj:`BufferError` is raised if neither is 

140 supported. 

141  

142 Since either way would take a consumer stream, for DLPack it is passed to 

143 ``obj.__dlpack__()`` as-is (except for :obj:`None`, see below); for CAI, a 

144 stream order will be established between the consumer stream and the 

145 producer stream (from ``obj.__cuda_array_interface__()["stream"]``), as if 

146 ``cudaStreamWaitEvent`` is called by this method. 

147  

148 To opt-out of the stream ordering operation in either DLPack or CAI, 

149 please pass ``stream_ptr=-1``. Note that this deviates (on purpose) 

150 from the semantics of ``obj.__dlpack__(stream=None, ...)`` since ``cuda.core`` 

151 does not encourage using the (legacy) default/null stream, but is 

152 consistent with the CAI's semantics. For DLPack, ``stream=-1`` will be 

153 internally passed to ``obj.__dlpack__()`` instead. 

154  

155 Parameters 

156 ---------- 

157 obj : Any 

158 Any objects that supports either DLPack (up to v1.0) or CUDA Array 

159 Interface (v3). 

160 stream_ptr: int 

161 The pointer address (as Python `int`) to the **consumer** stream. 

162 Stream ordering will be properly established unless ``-1`` is passed. 

163  

164  

165 Attributes 

166 ----------- 

167 ptr : int 

168 Pointer to the tensor buffer (as a Python `int`). 

169 device_id : int 

170 The device ID for where the tensor is located. It is -1 for CPU tensors 

171 (meaning those only accessible from the host). 

172 is_device_accessible : bool 

173 Whether the tensor data can be accessed on the GPU. 

174 readonly: bool 

175 Whether the tensor data can be modified in place. 

176 exporting_obj : Any 

177 A reference to the original tensor object that is being viewed. 

178 If the view is created with :meth:`from_buffer`, 

179 it will be the Buffer instance passed to the method. 

180  

181 """ 

182 def __init__(self, obj: object = None, stream_ptr: int | None = None) -> None: 

183 cdef str clsname = self.__class__.__name__ 7 ctx2W X Y Z 0 Ob}

184 if obj is not None: 7 ctx2W X Y Z 0 Ob}

185 # populate self's attributes 

186 if check_has_dlpack(obj): 6 ctx1WXYZ0}

187 warnings.warn( 5 ctx1WXYZ0

188 f"Constructing a {clsname} directly from a DLPack-supporting object is deprecated; " 5 ctx1WXYZ0

189 "Use `StridedMemoryView.from_dlpack` or `StridedMemoryView.from_any_interface` instead.", 

190 DeprecationWarning, 5 ctx1WXYZ0

191 stacklevel=2, 

192 ) 

193 view_as_dlpack(obj, stream_ptr, self) 5 ctx1WXYZ0

194 else: 

195 warnings.warn( 1 ctx1}

196 f"Constructing a {clsname} directly from a CUDA-array-interface-supporting object is deprecated; " 1 ctx1}

197 "Use `StridedMemoryView.from_cuda_array_interface` or `StridedMemoryView.from_any_interface` instead.", 

198 DeprecationWarning, 1 ctx1}

199 stacklevel=2, 

200 ) 

201 view_as_cai(obj, stream_ptr, self) 2 ctx1z}

202 else: 

203 warnings.warn( 1 ctx2Ob

204 f"Constructing an empty {clsname} is deprecated; " 1 ctx2Ob

205 "use one of the classmethods `from_dlpack`, `from_cuda_array_interface` or `from_any_interface` " 

206 "to construct a StridedMemoryView from an object", 

207 DeprecationWarning, 1 ctx2Ob

208 stacklevel=2, 

209 ) 

210  

211 @classmethod (empty)

212 def from_dlpack(cls, obj: object, stream_ptr: int | None=None) -> StridedMemoryView: 

213 """Create a view from an object supporting the `DLPack <https://dmlc.github.io/dlpack/latest/>`_ protocol. 

214  

215 Parameters 

216 ---------- 

217 obj : object 

218 An object implementing the `DLPack <https://dmlc.github.io/dlpack/latest/>`_ protocol 

219 (via ``__dlpack__``). 

220 stream_ptr : int, optional 

221 Stream pointer for synchronization. If ``None``, no synchronization is performed. 

222 """ 

223 cdef StridedMemoryView buf = StridedMemoryView.__new__(cls) 55 ctx2z 1 2 3 8 V P J K L M N Q R S T U x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

224 if _is_torch_tensor(obj): 54 ctx21 2 3 8 V P J K L M N Q R S T U x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

225 _get_tensor_bridge().view_as_torch_tensor(obj, stream_ptr, buf) 

226 return buf 

227 view_as_dlpack(obj, stream_ptr, buf) 54 ctx21 2 3 8 V P J K L M N Q R S T U x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

228 return buf 52 ctx11238VPJKLMNQRSTUxyw67cI5HBCDEFGAbsefhijopqrklmngutda

229  

230 @classmethod (empty)

231 def from_cuda_array_interface(cls, obj: object, stream_ptr: int | None=None) -> StridedMemoryView: 

232 """Create a view from an object supporting the `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_ protocol. 

233  

234 Parameters 

235 ---------- 

236 obj : object 

237 An object implementing the `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_ protocol. 

238 stream_ptr : int, optional 

239 Stream pointer for synchronization. If ``None``, no synchronization is performed. 

240 """ 

241 cdef StridedMemoryView buf = StridedMemoryView.__new__(cls) 6 ctx2z v | Ib~ 4

242 if _is_torch_tensor(obj): 5 ctx2v | Ib~ 4

243 _get_tensor_bridge().view_as_torch_tensor(obj, stream_ptr, buf) 

244 return buf 

245 view_as_cai(obj, stream_ptr, buf) 5 ctx2v | Ib~ 4

246 return buf 5 ctx1zv|~4

247  

248 @classmethod (empty)

249 def from_array_interface(cls, obj: object) -> StridedMemoryView: 

250 """Create a view from an object supporting the `__array_interface__ <https://numpy.org/doc/stable/reference/arrays.interface.html>`_ protocol. 

251  

252 Parameters 

253 ---------- 

254 obj : object 

255 An object implementing the `__array_interface__ <https://numpy.org/doc/stable/reference/arrays.interface.html>`_ protocol (e.g., a numpy array). 

256 """ 

257 cdef StridedMemoryView buf = StridedMemoryView.__new__(cls) 25 ctx2! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb

258 if _is_torch_tensor(obj): 26 ctx2z ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb

259 _get_tensor_bridge().view_as_torch_tensor(obj, None, buf) 

260 return buf 

261 view_as_array_interface(obj, buf) 25 ctx2! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb

262 return buf 24 ctx1!#-$.%'()*/:;=?@[]^+_`{,

263  

264 @classmethod (empty)

265 def from_any_interface(cls, obj: object, stream_ptr: int | None = None) -> StridedMemoryView: 

266 """Create a view by automatically selecting the best available protocol. 

267  

268 Tries `DLPack <https://dmlc.github.io/dlpack/latest/>`_ first, then falls back to 

269 `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_. 

270 ``torch.Tensor`` objects are transparently handled via a fast AOTI path 

271 regardless of which protocol is selected. 

272  

273 Parameters 

274 ---------- 

275 obj : object 

276 An object implementing `DLPack <https://dmlc.github.io/dlpack/latest/>`_ or 

277 `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_. 

278 stream_ptr : int, optional 

279 Stream pointer for synchronization. If ``None``, no synchronization is performed. 

280 """ 

281 if check_has_dlpack(obj): 42 ctx21 2 3 8 V P Q R S T U Zbx y w 6 7 c I 5 H A b s e f h i j o p q r k l m n g u t d a

282 return cls.from_dlpack(obj, stream_ptr) 41 ctx11238VPQRSTUxyw67cI5HAbsefhijopqrklmngutda

283 return cls.from_cuda_array_interface(obj, stream_ptr) 

284  

285 @classmethod (empty)

286 def from_buffer( 

287 cls, 

288 buffer : Buffer, 

289 shape : tuple[int, ...], 

290 strides : tuple[int, ...] | None = None, 

291 *, 

292 itemsize : int | None = None, (empty)

293 dtype : numpy.dtype | None = None, (empty)

294 is_readonly : bool = False (empty)

295 ) -> StridedMemoryView: 

296 """ 

297 Creates a :obj:`StridedMemoryView` instance from a :obj:`~_memory.Buffer` and shape and strides tuples. 

298 The Buffer can be either allocation coming from a :obj:`MemoryResource` or an external allocation 

299 wrapped in a :obj:`~_memory.Buffer` object with ``Buffer.from_handle(ptr, size, owner=...)``. 

300  

301 .. caution:: 

302 When creating a :obj:`StridedMemoryView` from a :obj:`~_memory.Buffer`, 

303 no synchronization is performed. It is the user's responsibility to ensure 

304 the data in ``buffer`` is properly synchronized when consuming the view. 

305  

306 Parameters 

307 ---------- 

308 buffer : :obj:`~_memory.Buffer` 

309 The buffer to create the view from. 

310 shape : :obj:`tuple` 

311 The layout describing the shape, strides and itemsize of the elements in 

312 the buffer. 

313 strides : :obj:`tuple` 

314 The layout describing the shape, strides and itemsize of the elements in 

315 the buffer. 

316 dtype : :obj:`numpy.dtype` 

317 Optional dtype. 

318 If specified, the dtype's itemsize must match the layout's itemsize. 

319 is_readonly : bool, optional 

320 Whether the mark the view as readonly. 

321 """ 

322 cdef StridedMemoryView view = StridedMemoryView.__new__(cls) 43 ctx2RbSbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbUbQbabbbNbJbO 9

323 if itemsize is None and dtype is None: 43 ctx2RbSbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbUbQbabbbNbJbO 9

324 raise ValueError("Either itemsize or dtype must be specified") 1 ctx2Ub

325 if itemsize is not None and dtype is not None and itemsize != dtype.itemsize: 42 ctx2RbSbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbQbabbbNbJbO 9

326 raise ValueError( 1 ctx2Qb

327 f"itemsize ({itemsize}) does not match dtype.itemsize ({dtype.itemsize})" 1 ctx2Qb

328 ) 

329 # (itemsize is None XOR dtype is None) OR they are equal 

330 view_buffer_strided( 39 ctx2cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbabbbNbJbO 9

331 view, 

332 buffer, 

333 _StridedLayout(shape=shape, strides=strides, itemsize=getattr(dtype, "itemsize", itemsize)), 41 ctx2RbSbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbabbbNbJbO 9

334 dtype, 

335 is_readonly, 

336 ) 

337 return view 37 ctx2cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbbJbO 9

338  

339 def __dealloc__(self) -> None: 

340 if self.dl_tensor == NULL: 133 ctx2RbSb1 2 3 8 V P J K L M N Q R S T U W X Y Z 0 v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , KbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbUbQbabbbNbJb| Ib~ 6 7 c O Ob} 9 I 5 H 4 B C D E F G A s e f h i j o p q r k l m n g u t d Lba Mb

341 return 88 ctx2RbSbv x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , KbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbPbUbQbabbbNbJb| Ib~ O Ob} 9 5 H 4 B C D E F G LbMb

342  

343 if cpython.PyCapsule_IsValid( 56 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAsefhijopqrklmngutda

344 self.metadata, DLPACK_VERSIONED_TENSOR_USED_NAME): 56 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAsefhijopqrklmngutda

345 data = cpython.PyCapsule_GetPointer( 56 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAsefhijopqrklmngutda

346 self.metadata, DLPACK_VERSIONED_TENSOR_USED_NAME) 56 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAsefhijopqrklmngutda

347 dlm_tensor_ver = <DLManagedTensorVersioned*>data 56 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAsefhijopqrklmngutda

348 dlm_tensor_ver.deleter(dlm_tensor_ver) 56 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAsefhijopqrklmngutda

349 elif cpython.PyCapsule_IsValid( 1 ctx1a

350 self.metadata, DLPACK_TENSOR_USED_NAME): 1 ctx1a

351 data = cpython.PyCapsule_GetPointer( 2 ctx1za

352 self.metadata, DLPACK_TENSOR_USED_NAME) 1 ctx1a

353 dlm_tensor = <DLManagedTensor*>data 1 ctx1a

354 dlm_tensor.deleter(dlm_tensor) 1 ctx1a

355  

356 def view( (empty)

357 self, layout : _StridedLayout | None = None, dtype : numpy.dtype | None = None 

358 ) -> StridedMemoryView: 

359 """ 

360 Creates a new view with adjusted layout and dtype. 

361 Same as calling :meth:`from_buffer` with the current buffer. 

362 """ 

363 cdef StridedMemoryView view = StridedMemoryView.__new__(self.__class__) 13 ctx2x y w abbb5 H B C D E F G

364 if layout is None and dtype is None: 13 ctx2x y w abbb5 H B C D E F G

365 return self 2 ctx1z5

366 if layout is None: 12 ctx2x y w abbbH B C D E F G

367 layout = self.get_layout() 4 ctx1xywH

368 if dtype is None: 12 ctx2x y w abbbH B C D E F G

369 dtype = self.get_dtype() 8 ctx2abbbB C D E F G

370 view_buffer_strided(view, self.get_buffer(), layout, dtype, self.readonly) 13 ctx2z x y w abbbH B C D E F G

371 return view 12 ctx2x y w abbbH B C D E F G

372  

373 def as_tensor_map( (empty)

374 self, 

375 box_dim: tuple[int, ...] | None = None, 

376 *, 

377 options: object = None, (empty)

378 element_strides: tuple[int, ...] | None = None, (empty)

379 data_type: object = None, (empty)

380 interleave: object = None, (empty)

381 swizzle: object = None, (empty)

382 l2_promotion: object = None, (empty)

383 oob_fill: object = None, (empty)

384 ) -> object: 

385 """Create a tiled :obj:`TensorMapDescriptor` from this view. 

386  

387 This is the public entry point for creating tiled tensor map 

388 descriptors in ``cuda.core``. Pass either ``box_dim`` and the 

389 individual keyword arguments directly, or provide bundled tiled 

390 options via ``options=``. 

391 """ 

392 from cuda.core._tensor_map import TensorMapDescriptor 1 ctx1P

393  

394 kwargs = {} 1 ctx1P

395 if options is not None: 1 ctx1P

396 kwargs["options"] = options 

397 if element_strides is not None: 1 ctx1P

398 kwargs["element_strides"] = element_strides 1 ctx1P

399 if data_type is not None: 1 ctx1P

400 kwargs["data_type"] = data_type 1 ctx1P

401 if interleave is not None: 2 ctx1zP

402 kwargs["interleave"] = interleave 

403 if swizzle is not None: 1 ctx1P

404 kwargs["swizzle"] = swizzle 1 ctx1P

405 if l2_promotion is not None: 1 ctx1P

406 kwargs["l2_promotion"] = l2_promotion 1 ctx1P

407 if oob_fill is not None: 1 ctx1P

408 kwargs["oob_fill"] = oob_fill 1 ctx1P

409 return TensorMapDescriptor._from_tiled(self, box_dim, **kwargs) 1 ctx1P

410  

411 def copy_from( (empty)

412 self, 

413 other: StridedMemoryView, 

414 stream: Stream, 

415 allocator: object = None, 

416 blocking: bool | None = None, 

417 ) -> None: 

418 """ 

419 Copies the data from the other view into this view. 

420  

421 The copy can be performed between following memory spaces: 

422 host-to-device, device-to-host, device-to-device (on the same device). 

423  

424 Parameters 

425 ---------- 

426 other : StridedMemoryView 

427 The view to copy data from. 

428 stream : Stream | None, optional 

429 The stream to schedule the copy on. 

430 allocator : MemoryResource | None, optional 

431 If temporary buffers are needed, the specified memory resources 

432 will be used to allocate the memory. If not specified, default 

433 resources will be used. 

434 blocking : bool | None, optional 

435 Whether the call should block until the copy is complete. 

436 * ``True``: the ``stream`` is synchronized with the host at the end of the call, 

437 blocking until the copy is complete. 

438 * ``False``: if possible, the call returns immediately once the copy is scheduled. 

439 However, in some cases of host-to-device or device-to-host copies, the call may 

440 still synchronize with the host if necessary. 

441 * ``None`` (default): 

442 * for device-to-device, it defaults to ``False`` (non-blocking), 

443 * for host-to-device or device-to-host, it defaults to ``True`` (blocking). 

444 """ 

445 raise NotImplementedError("Sorry, not supported: copy_from") 1 ctx16

446  

447 def copy_to( (empty)

448 self, 

449 other: StridedMemoryView, 

450 stream: Stream | None = None, 

451 allocator: object = None, 

452 blocking: bool | None = None, 

453 ) -> None: 

454 """ 

455 Copies the data from this view into the ``other`` view. 

456  

457 For details, see :meth:`copy_from`. 

458 """ 

459 raise NotImplementedError("Sorry, not supported: copy_to") 1 ctx17

460  

461 def __dlpack__( (empty)

462 self, 

463 *, 

464 stream: int | None = None, (empty)

465 max_version: tuple[int, int] | None = None, (empty)

466 dl_device: tuple[int, int] | None = None, (empty)

467 copy: bool | None = None, (empty)

468 ) -> object: 

469 # Similar to Buffer.__dlpack__: no implicit synchronization is performed. 

470 if dl_device is not None: 26 ctx13vxywcOsefhijopqrklmngutda

471 raise BufferError("Sorry, not supported: dl_device other than None") 1 ctx13

472 if copy is True: 26 ctx13vxywcOsefhijopqrklmngutda

473 raise BufferError("Sorry, not supported: copy=True") 1 ctx13

474  

475 cdef bint versioned 

476 if max_version is None: 26 ctx13vxywcOsefhijopqrklmngutda

477 versioned = False 7 ctx1vxywOta

478 else: 

479 if not isinstance(max_version, tuple) or len(max_version) != 2: 20 ctx13csefhijopqrklmnguda

480 raise BufferError(f"Expected max_version tuple[int, int], got {max_version}") 1 ctx13

481 versioned = max_version >= (1, 0) 20 ctx1zcsefhijopqrklmnguda

482  

483 # NOTE: stream is accepted for protocol compatibility but not used. 

484 cdef object capsule = _smv_make_py_capsule(self, versioned) 25 ctx1vxywcOsefhijopqrklmngutda

485 return capsule 21 ctx1vcsefhijopqrklmngutda

486  

487 def __dlpack_device__(self) -> tuple[int, int]: (empty)

488 cdef _DLDeviceType device_type 

489 cdef int32_t device_id 

490 _smv_get_dl_device(self, &device_type, &device_id) 5 ctx112vca

491 return (<int>device_type, int(device_id)) 5 ctx112vca

492  

493 @property 

494 def _layout(self) -> _StridedLayout: 

495 """ 

496 The layout of the tensor. For StridedMemoryView created from DLPack or CAI, 

497 the layout is inferred from the tensor object's metadata. 

498 """ 

499 return self.get_layout() 42 ctx2cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbbObH B C D E F G

500  

501 @property 

502 def size(self) -> int: 

503 return self.get_layout().get_volume() 40 ctx1JKLMNQRSTUWXYZ0!#-$.%'()*/:;=?@[]^+_`{,|

504  

505 @property 

506 def shape(self) -> tuple[int, ...]: 

507 """ 

508 Shape of the tensor. 

509 """ 

510 return self.get_layout().get_shape_tuple() 87 ctx2V J K L M N Q R S T U W X Y Z 0 ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbb| ~ } 9 I 4 B C D E F G b

511  

512 @property 

513 def strides(self) -> tuple[int, ...] | None: 

514 """ 

515 Strides of the tensor (in **counts**, not bytes). 

516 """ 

517 return self.get_layout().get_strides_tuple() 76 ctx2V J K L M N Q R S T U W X Y Z 0 ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHb| ~ 9 I

518  

519 @property 

520 def dtype(self) -> numpy.dtype | None: 

521 """ 

522 Data type of the tensor. 

523  

524 Supports standard NumPy dtypes as well as narrow data types (e.g., ``bfloat16``) 

525 when the optional `ml_dtypes <https://github.com/jax-ml/ml_dtypes>`_ package is 

526 installed. If ``ml_dtypes`` is not available and such a tensor is encountered, 

527 a :obj:`NotImplementedError` will be raised. 

528 """ 

529 return self.get_dtype() 66 ctx2! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbJb9 I H B C D E F G

530  

531 def __repr__(self) -> str: 

532 return (f"StridedMemoryView(ptr={self.ptr},\n" 2 ctx19I

533 + f" shape={self.shape},\n" 2 ctx19I

534 + f" strides={self.strides},\n" 2 ctx19I

535 + f" itemsize={self._layout.itemsize},\n" 2 ctx19I

536 + f" dtype={get_simple_repr(self.dtype)},\n" 2 ctx19I

537 + f" device_id={self.device_id},\n" 2 ctx19I

538 + f" is_device_accessible={self.is_device_accessible},\n" 2 ctx19I

539 + f" readonly={self.readonly},\n" 2 ctx19I

540 + f" exporting_obj={get_simple_repr(self.exporting_obj)})") 2 ctx19I

541  

542 cdef inline _StridedLayout get_layout(self): (empty)

543 if self._layout is None: 117 ctx2V J K L M N Q R S T U W X Y Z 0 v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , KbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbb| Ib~ c O Ob} 9 I H 4 B C D E F G A b s e f h i j o p q r k l m n g u t d a

544 if self.dl_tensor: 81 ctx2V J K L M N Q R S T U W X Y Z 0 v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb| Ib~ c Ob} I H 4 B C D E F G A b s e f h i j o p q r k l m n g u t d a

545 self._layout = layout_from_dlpack(self.dl_tensor) 49 ctx1VJKLMNQRSTUWXYZ0xywcIHBCDEFGAbsefhijopqrklmngutda

546 elif self.metadata is not None: 32 ctx2v ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb| Ib~ Ob} 4

547 self._layout = layout_from_cai(self.metadata) 31 ctx2v ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , Kb| Ib~ } 4

548 else: 

549 raise ValueError("Cannot infer layout from the exporting object") 1 ctx2Ob

550 return self._layout 114 ctx2V J K L M N Q R S T U W X Y Z 0 v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbb| ~ c O } 9 I H 4 B C D E F G A b s e f h i j o p q r k l m n g u t d a

551  

552 cdef inline object get_buffer(self): (empty)

553 """ 

554 Returns Buffer instance with the underlying data. 

555 If the SMV was created from a Buffer, it will return the same Buffer instance. 

556 Otherwise, it will create a new instance with owner set to the exporting object. 

557 """ 

558 if self._buffer is None: 13 ctx2v x y w abbbH B C D E F G

559 if isinstance(self.exporting_obj, Buffer): 11 ctx1vxywHBCDEFG

560 self._buffer = self.exporting_obj 

561 else: 

562 self._buffer = Buffer.from_handle(self.ptr, 0, owner=self.exporting_obj) 11 ctx1vxywHBCDEFG

563 return self._buffer 13 ctx2v x y w abbbH B C D E F G

564  

565 cdef inline object get_dtype(self): (empty)

566 if self._dtype is None: 95 ctx2v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbbJbc O 9 I H B C D E F G A b s e f h i j o p q r k l m n g u t d a

567 if self.dl_tensor != NULL: 56 ctx1v!#-$.%'()*/:;=?@[]^+_`{,cO9IBCDEFGAbsefhijopqrklmngutda

568 self._dtype = dtype_dlpack_to_numpy(&self.dl_tensor.dtype) 29 ctx1cIBCDEFGAbsefhijopqrklmngutda

569 elif isinstance(self.metadata, int): 27 ctx1v!#-$.%'()*/:;=?@[]^+_`{,O9

570 # AOTI dtype code stored by the torch tensor bridge 

571 self._dtype = _get_tensor_bridge().resolve_aoti_dtype( 

572 self.metadata) 

573 elif self.metadata is not None: 27 ctx1v!#-$.%'()*/:;=?@[]^+_`{,O9

574 self._dtype = _typestr2dtype(self.metadata["typestr"]) 25 ctx1v!#-$.%'()*/:;=?@[]^+_`{,

575 return self._dtype 95 ctx2v x y w ! # - $ . % ' ( ) * / : ; = ? @ [ ] ^ + _ ` { , cbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbabbbJbc O 9 I H B C D E F G A b s e f h i j o p q r k l m n g u t d a

576  

577  

578cdef void _smv_pycapsule_deleter(object capsule) noexcept: (empty)

579 cdef DLManagedTensor* dlm_tensor 

580 cdef DLManagedTensorVersioned* dlm_tensor_ver 

581 # Do not invoke the deleter on a used capsule. 

582 if cpython.PyCapsule_IsValid(capsule, DLPACK_TENSOR_UNUSED_NAME): 21 ctx1vcsefhijopqrklmngutda

583 dlm_tensor = <DLManagedTensor*>( 3 ctx1zvt

584 cpython.PyCapsule_GetPointer(capsule, DLPACK_TENSOR_UNUSED_NAME) 2 ctx1vt

585 ) 

586 if dlm_tensor.deleter: 2 ctx1vt

587 dlm_tensor.deleter(dlm_tensor) 3 ctx1zvt

588 elif cpython.PyCapsule_IsValid(capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME): 19 ctx1csefhijopqrklmnguda

589 dlm_tensor_ver = <DLManagedTensorVersioned*>( 1 ctx1d

590 cpython.PyCapsule_GetPointer(capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME) 1 ctx1d

591 ) 

592 if dlm_tensor_ver.deleter: 1 ctx1d

593 dlm_tensor_ver.deleter(dlm_tensor_ver) 1 ctx1d

594  

595  

596cdef inline void _smv_release_export_resources(void* manager_ctx, int64_t* shape_ptr) noexcept with gil: (empty)

597 if shape_ptr: 25 ctx1vxywcOsefhijopqrklmngutda

598 stdlib.free(shape_ptr) 20 ctx1vcsefhijopqrklmngtda

599 if manager_ctx: 25 ctx1vxywcOsefhijopqrklmngutda

600 cpython.Py_DECREF(<object>manager_ctx) 25 ctx1vxywcOsefhijopqrklmngutda

601  

602  

603cdef void _smv_deleter(DLManagedTensor* tensor) noexcept with gil: (empty)

604 if tensor: 7 ctx1vxywOta

605 _smv_release_export_resources(tensor.manager_ctx, tensor.dl_tensor.shape) 7 ctx1vxywOta

606 tensor.manager_ctx = NULL 7 ctx1vxywOta

607 stdlib.free(tensor) 7 ctx1vxywOta

608  

609  

610cdef void _smv_versioned_deleter(DLManagedTensorVersioned* tensor) noexcept with gil: (empty)

611 if tensor: 23 ctx1xywcOsefhijopqrklmnguda

612 _smv_release_export_resources(tensor.manager_ctx, tensor.dl_tensor.shape) 19 ctx1csefhijopqrklmnguda

613 tensor.manager_ctx = NULL 20 ctx1zcsefhijopqrklmnguda

614 stdlib.free(tensor) 19 ctx1csefhijopqrklmnguda

615  

616  

617cdef inline DLManagedTensorVersioned* _smv_allocate_dlm_tensor_versioned() except? NULL: (empty)

618 cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL 20 ctx1cbsefhijopqrklmnguda

619 dlm_tensor_ver = <DLManagedTensorVersioned*>stdlib.malloc(sizeof(DLManagedTensorVersioned)) 20 ctx1cbsefhijopqrklmnguda

620 if dlm_tensor_ver == NULL: 20 ctx1cbsefhijopqrklmnguda

621 raise MemoryError() (empty)

622 dlm_tensor_ver.dl_tensor.shape = NULL 20 ctx1cbsefhijopqrklmnguda

623 dlm_tensor_ver.manager_ctx = NULL 20 ctx1cbsefhijopqrklmnguda

624 return dlm_tensor_ver 20 ctx1cbsefhijopqrklmnguda

625  

626  

627cdef inline DLManagedTensor* _smv_allocate_dlm_tensor() except? NULL: (empty)

628 cdef DLManagedTensor* dlm_tensor = NULL 7 ctx1vxywOta

629 dlm_tensor = <DLManagedTensor*>stdlib.malloc(sizeof(DLManagedTensor)) 7 ctx1vxywOta

630 if dlm_tensor == NULL: 7 ctx1vxywOta

631 raise MemoryError() 

632 dlm_tensor.dl_tensor.shape = NULL 8 ctx1zvxywOta

633 dlm_tensor.manager_ctx = NULL 7 ctx1vxywOta

634 return dlm_tensor 7 ctx1vxywOta

635  

636  

637cdef inline int _smv_dtype_numpy_to_dlpack(object dtype_obj, DLDataType* out_dtype) except -1: (empty)

638 cdef object np_dtype = numpy.dtype(dtype_obj) 26 ctx1vxywcAbsefhijopqrklmngutda

639 if np_dtype.fields is not None: 26 ctx1vxywcAbsefhijopqrklmngutda

640 raise BufferError("Structured dtypes are not supported for DLPack export") 2 ctx1zy

641 if not np_dtype.isnative and np_dtype.byteorder not in ("=", "|"): 25 ctx1vxwcAbsefhijopqrklmngutda

642 raise BufferError("Non-native-endian dtypes are not supported for DLPack export") 1 ctx1x

643  

644 cdef str kind = np_dtype.kind 24 ctx1vwcAbsefhijopqrklmngutda

645 cdef int bits = np_dtype.itemsize * 8 24 ctx1vwcAbsefhijopqrklmngutda

646 cdef uint8_t code 

647 if kind == "b": 24 ctx1vwcAbsefhijopqrklmngutda

648 if bits != 8: 2 ctx1zs

649 raise BufferError(f"Unsupported bool dtype itemsize: {np_dtype.itemsize}") 

650 code = <uint8_t>kDLBool 1 ctx1s

651 elif kind == "i": 23 ctx1vwcAbefhijopqrklmngutda

652 if bits not in (8, 16, 32, 64): 9 ctx1cAopqrtda

653 raise BufferError(f"Unsupported signed integer dtype: {np_dtype}") 

654 code = <uint8_t>kDLInt 9 ctx1cAopqrtda

655 elif kind == "u": 14 ctx1vwbefhijklmngu

656 if bits not in (8, 16, 32, 64): 4 ctx1klmn

657 raise BufferError(f"Unsupported unsigned integer dtype: {np_dtype}") 

658 code = <uint8_t>kDLUInt 4 ctx1klmn

659 elif kind == "f": 10 ctx1vwbefhijgu

660 if bits not in (16, 32, 64): 5 ctx1vbhij

661 raise BufferError(f"Unsupported floating dtype: {np_dtype}") 

662 code = <uint8_t>kDLFloat 5 ctx1vbhij

663 elif kind == "c": 5 ctx1wefgu

664 if bits not in (64, 128): 4 ctx1efgu

665 raise BufferError(f"Unsupported complex dtype: {np_dtype}") 

666 code = <uint8_t>kDLComplex 4 ctx1efgu

667 else: 

668 raise BufferError(f"Unsupported dtype for DLPack export: {np_dtype}") 2 ctx1zw

669  

670 out_dtype.code = code 23 ctx1vcAbsefhijopqrklmngutda

671 out_dtype.bits = <uint8_t>bits 23 ctx1vcAbsefhijopqrklmngutda

672 out_dtype.lanes = <uint16_t>1 23 ctx1vcAbsefhijopqrklmngutda

673 return 0 24 ctx1zvcAbsefhijopqrklmngutda

674  

675  

676cdef inline int _smv_get_dl_device( (empty)

677 StridedMemoryView view, 

678 _DLDeviceType* out_device_type, 

679 int32_t* out_device_id, 

680) except -1: 

681 cdef _DLDeviceType device_type 

682 cdef int32_t device_id 

683 cdef object buf 

684 if view.dl_tensor != NULL: 25 ctx112vcAbsefhijopqrklmngutda

685 device_type = view.dl_tensor.device.device_type 24 ctx112cAbsefhijopqrklmngutda

686 if device_type == _kDLCUDA: 24 ctx112cAbsefhijopqrklmngutda

687 device_id = view.dl_tensor.device.device_id 

688 else: 

689 # CPU, CUDAHost, and CUDAManaged use device_id=0 in DLPack. 

690 device_id = 0 24 ctx112cAbsefhijopqrklmngutda

691 elif view.is_device_accessible: 1 ctx1v

692 buf = view.get_buffer() 1 ctx1v

693 dev_type, dev_id = classify_dl_device(buf) 1 ctx1v

694 device_type = <_DLDeviceType>dev_type 1 ctx1v

695 device_id = <int32_t>dev_id 1 ctx1v

696 else: 

697 device_type = _kDLCPU 

698 device_id = 0 

699  

700 out_device_type[0] = device_type 25 ctx112vcAbsefhijopqrklmngutda

701 out_device_id[0] = device_id 25 ctx112vcAbsefhijopqrklmngutda

702 return 0 25 ctx112vcAbsefhijopqrklmngutda

703  

704  

705cdef inline int _smv_setup_dl_tensor_common( (empty)

706 DLTensor* dl_tensor, 

707 StridedMemoryView view, 

708 _StridedLayout layout, 

709) except -1: 

710 cdef object dtype_obj = view.get_dtype() 27 ctx1vxywcOAbsefhijopqrklmngutda

711 if dtype_obj is None: 27 ctx1vxywcOAbsefhijopqrklmngutda

712 raise BufferError( 1 ctx1O

713 "Cannot export StridedMemoryView via DLPack without dtype information; " 

714 "create the view with dtype specified." 

715 ) 

716 _smv_dtype_numpy_to_dlpack(dtype_obj, &dl_tensor.dtype) 26 ctx1vxywcAbsefhijopqrklmngutda

717 _smv_get_dl_device(view, &dl_tensor.device.device_type, &dl_tensor.device.device_id) 23 ctx1vcAbsefhijopqrklmngutda

718  

719 cdef int ndim = layout.base.ndim 23 ctx1vcAbsefhijopqrklmngutda

720 dl_tensor.ndim = ndim 23 ctx1vcAbsefhijopqrklmngutda

721 if layout.get_volume() == 0: 23 ctx1vcAbsefhijopqrklmngutda

722 dl_tensor.data = NULL 1 ctx1g

723 else: 

724 dl_tensor.data = <void*><intptr_t>view.ptr 22 ctx1vcAbsefhijopqrklmnutda

725 dl_tensor.byte_offset = 0 23 ctx1vcAbsefhijopqrklmngutda

726 return 0 23 ctx1vcAbsefhijopqrklmngutda

727  

728  

729cdef inline int _smv_setup_dl_tensor(DLTensor* dl_tensor, StridedMemoryView view) except -1: (empty)

730 cdef _StridedLayout layout = view.get_layout() 26 ctx1vxywcObsefhijopqrklmngutda

731 _smv_setup_dl_tensor_common(dl_tensor, view, layout) 26 ctx1vxywcObsefhijopqrklmngutda

732  

733 cdef int i 

734 cdef int64_t* shape_strides = NULL 22 ctx1vcbsefhijopqrklmngutda

735 cdef int64_t* strides_src = NULL 22 ctx1vcbsefhijopqrklmngutda

736 cdef int ndim = dl_tensor.ndim 22 ctx1vcbsefhijopqrklmngutda

737 if ndim == 0: 22 ctx1vcbsefhijopqrklmngutda

738 dl_tensor.shape = NULL 1 ctx1u

739 dl_tensor.strides = NULL 1 ctx1u

740 else: 

741 # DLPack v1.2+ requires non-NULL strides for ndim != 0. 

742 shape_strides = <int64_t*>stdlib.malloc(sizeof(int64_t) * 2 * ndim) 21 ctx1vcbsefhijopqrklmngtda

743 if shape_strides == NULL: 21 ctx1vcbsefhijopqrklmngtda

744 raise MemoryError() 

745 try: 21 ctx1vcbsefhijopqrklmngtda

746 strides_src = get_strides_ptr(layout.base) 21 ctx1vcbsefhijopqrklmngtda

747 for i in range(ndim): 21 ctx1vcbsefhijopqrklmngtda

748 shape_strides[i] = layout.base.shape[i] 21 ctx1vcbsefhijopqrklmngtda

749 shape_strides[i + ndim] = strides_src[i] 21 ctx1vcbsefhijopqrklmngtda

750 except Exception: 

751 stdlib.free(shape_strides) 

752 raise 

753 dl_tensor.shape = shape_strides 21 ctx1vcbsefhijopqrklmngtda

754 dl_tensor.strides = shape_strides + ndim 21 ctx1vcbsefhijopqrklmngtda

755 return 0 22 ctx1vcbsefhijopqrklmngutda

756  

757  

758cdef inline int _smv_setup_dltensor_borrowed(DLTensor* dl_tensor, StridedMemoryView view) except -1: (empty)

759 cdef _StridedLayout layout = view.get_layout() 1 ctx1A

760 _smv_setup_dl_tensor_common(dl_tensor, view, layout) 1 ctx1A

761  

762 if dl_tensor.ndim == 0: 1 ctx1A

763 dl_tensor.shape = NULL 

764 dl_tensor.strides = NULL 

765 else: 

766 dl_tensor.shape = layout.base.shape 1 ctx1A

767 # For temporary/non-owning exchange we provide explicit strides. 

768 dl_tensor.strides = get_strides_ptr(layout.base) 1 ctx1A

769 return 0 1 ctx1A

770  

771  

772cdef inline int _smv_fill_managed_tensor_versioned( (empty)

773 DLManagedTensorVersioned* dlm_tensor_ver, 

774 StridedMemoryView view, 

775) except -1: 

776 cpython.Py_INCREF(view) 20 ctx1cbsefhijopqrklmnguda

777 dlm_tensor_ver.manager_ctx = <void*>view 20 ctx1cbsefhijopqrklmnguda

778 dlm_tensor_ver.deleter = _smv_versioned_deleter 20 ctx1cbsefhijopqrklmnguda

779 dlm_tensor_ver.version.major = DLPACK_MAJOR_VERSION 20 ctx1cbsefhijopqrklmnguda

780 dlm_tensor_ver.version.minor = DLPACK_MINOR_VERSION 20 ctx1cbsefhijopqrklmnguda

781 dlm_tensor_ver.flags = DLPACK_FLAG_BITMASK_READ_ONLY if view.readonly else 0 20 ctx1cbsefhijopqrklmnguda

782 _smv_setup_dl_tensor(&dlm_tensor_ver.dl_tensor, view) 20 ctx1cbsefhijopqrklmnguda

783 return 0 20 ctx1cbsefhijopqrklmnguda

784  

785  

786cdef inline int _smv_fill_managed_tensor( (empty)

787 DLManagedTensor* dlm_tensor, 

788 StridedMemoryView view, 

789) except -1: 

790 cpython.Py_INCREF(view) 7 ctx1vxywOta

791 dlm_tensor.manager_ctx = <void*>view 7 ctx1vxywOta

792 dlm_tensor.deleter = _smv_deleter 7 ctx1vxywOta

793 _smv_setup_dl_tensor(&dlm_tensor.dl_tensor, view) 7 ctx1vxywOta

794 return 0 3 ctx1vta

795  

796  

797cdef object _smv_make_py_capsule(StridedMemoryView view, bint versioned): (empty)

798 cdef DLManagedTensor* dlm_tensor = NULL 25 ctx1vxywcOsefhijopqrklmngutda

799 cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL 25 ctx1vxywcOsefhijopqrklmngutda

800 cdef object capsule = None 25 ctx1vxywcOsefhijopqrklmngutda

801 cdef void* tensor_ptr = NULL 25 ctx1vxywcOsefhijopqrklmngutda

802 cdef const char* capsule_name 

803 try: 25 ctx1vxywcOsefhijopqrklmngutda

804 if versioned: 25 ctx1vxywcOsefhijopqrklmngutda

805 dlm_tensor_ver = _smv_allocate_dlm_tensor_versioned() 19 ctx1csefhijopqrklmnguda

806 _smv_fill_managed_tensor_versioned(dlm_tensor_ver, view) 19 ctx1csefhijopqrklmnguda

807 tensor_ptr = <void*>dlm_tensor_ver 19 ctx1csefhijopqrklmnguda

808 capsule_name = DLPACK_VERSIONED_TENSOR_UNUSED_NAME 19 ctx1csefhijopqrklmnguda

809 else: 

810 dlm_tensor = _smv_allocate_dlm_tensor() 7 ctx1vxywOta

811 _smv_fill_managed_tensor(dlm_tensor, view) 7 ctx1vxywOta

812 tensor_ptr = <void*>dlm_tensor 3 ctx1vta

813 capsule_name = DLPACK_TENSOR_UNUSED_NAME 3 ctx1vta

814 capsule = cpython.PyCapsule_New(tensor_ptr, capsule_name, _smv_pycapsule_deleter) 21 ctx1vcsefhijopqrklmngutda

815 except Exception: 4 ctx1xywO

816 if capsule is None: 4 ctx1xywO

817 _smv_deleter(dlm_tensor) 4 ctx1xywO

818 _smv_versioned_deleter(dlm_tensor_ver) 4 ctx1xywO

819 raise 4 ctx1xywO

820 return capsule 21 ctx1vcsefhijopqrklmngutda

821  

822  

823cdef inline StridedMemoryView _smv_from_dlpack_capsule(object capsule, object exporting_obj): (empty)

824 cdef void* data = NULL 1 ctx1b

825 cdef DLTensor* dl_tensor = NULL 1 ctx1b

826 cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL 1 ctx1b

827 cdef DLManagedTensor* dlm_tensor = NULL 1 ctx1b

828 cdef bint is_readonly = False 1 ctx1b

829 cdef const char* used_name = NULL 1 ctx1b

830 if cpython.PyCapsule_IsValid(capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME): 1 ctx1b

831 data = cpython.PyCapsule_GetPointer(capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME) 1 ctx1b

832 dlm_tensor_ver = <DLManagedTensorVersioned*>data 1 ctx1b

833 dl_tensor = &dlm_tensor_ver.dl_tensor 1 ctx1b

834 is_readonly = bool((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0) 1 ctx1b

835 used_name = DLPACK_VERSIONED_TENSOR_USED_NAME 1 ctx1b

836 elif cpython.PyCapsule_IsValid(capsule, DLPACK_TENSOR_UNUSED_NAME): 

837 data = cpython.PyCapsule_GetPointer(capsule, DLPACK_TENSOR_UNUSED_NAME) 

838 dlm_tensor = <DLManagedTensor*>data 

839 dl_tensor = &dlm_tensor.dl_tensor 

840 is_readonly = False 

841 used_name = DLPACK_TENSOR_USED_NAME 

842 else: 

843 raise BufferError("Invalid DLPack capsule") 

844  

845 cpython.PyCapsule_SetName(capsule, used_name) 1 ctx1b

846  

847 cdef StridedMemoryView view = StridedMemoryView.__new__(StridedMemoryView) 1 ctx1b

848 view.dl_tensor = dl_tensor 1 ctx1b

849 view.metadata = capsule 1 ctx1b

850 view.ptr = <intptr_t>(dl_tensor.data) + <intptr_t>(dl_tensor.byte_offset) 1 ctx1b

851 view.readonly = is_readonly 1 ctx1b

852 view.exporting_obj = exporting_obj 1 ctx1b

853 if dl_tensor.device.device_type == _kDLCPU: 1 ctx1b

854 view.device_id = -1 1 ctx1b

855 view.is_device_accessible = False 1 ctx1b

856 elif dl_tensor.device.device_type in (_kDLCUDA, _kDLCUDAHost, _kDLCUDAManaged): 

857 view.device_id = dl_tensor.device.device_id 

858 view.is_device_accessible = True 

859 else: 

860 raise BufferError("device not supported") 

861 return view 1 ctx1b

862  

863  

864cdef int _smv_managed_tensor_allocator( (empty)

865 DLTensor* prototype, 

866 DLManagedTensorVersioned** out, 

867 void* error_ctx, 

868 void (*SetError)(void* error_ctx, const char* kind, const char* message) noexcept, 

869) noexcept with gil: 

870 if out != NULL: 1 ctx2Vb

871 out[0] = NULL 1 ctx2Vb

872 if SetError != NULL: 1 ctx2Vb

873 SetError(error_ctx, b"NotImplementedError", b"managed_tensor_allocator is not supported by StridedMemoryView") 

874 cpython.PyErr_SetString(NotImplementedError, b"managed_tensor_allocator is not supported by StridedMemoryView") 1 ctx2Vb

875 return -1 1 ctx2Vb

876  

877  

878cdef int _smv_managed_tensor_from_py_object_no_sync( (empty)

879 void* py_object, 

880 DLManagedTensorVersioned** out, 

881) noexcept with gil: 

882 cdef DLManagedTensorVersioned* dlm_tensor_ver = NULL 1 ctx1b

883 if out == NULL: 1 ctx1b

884 cpython.PyErr_SetString(RuntimeError, b"out cannot be NULL") 

885 return -1 

886 out[0] = NULL 1 ctx1b

887 cdef object obj = <object>py_object 1 ctx1b

888 if not isinstance(obj, StridedMemoryView): 1 ctx1b

889 cpython.PyErr_SetString(TypeError, b"py_object must be a StridedMemoryView") 

890 return -1 

891 try: 1 ctx1b

892 dlm_tensor_ver = _smv_allocate_dlm_tensor_versioned() 1 ctx1b

893 _smv_fill_managed_tensor_versioned(dlm_tensor_ver, <StridedMemoryView>obj) 1 ctx1b

894 except Exception: 

895 _smv_versioned_deleter(dlm_tensor_ver) 

896 return -1 

897 out[0] = dlm_tensor_ver 1 ctx1b

898 return 0 1 ctx1b

899  

900  

901cdef int _smv_managed_tensor_to_py_object_no_sync( (empty)

902 DLManagedTensorVersioned* tensor, 

903 void** out_py_object, 

904) noexcept with gil: 

905 cdef object capsule 

906 cdef object py_view 

907 if out_py_object == NULL: 2 ctx2b Wb

908 cpython.PyErr_SetString(RuntimeError, b"out_py_object cannot be NULL") 

909 return -1 

910 out_py_object[0] = NULL 2 ctx2b Wb

911 if tensor == NULL: 2 ctx2b Wb

912 cpython.PyErr_SetString(RuntimeError, b"tensor cannot be NULL") 1 ctx2Wb

913 return -1 1 ctx2Wb

914 try: 1 ctx1b

915 capsule = cpython.PyCapsule_New( 1 ctx1b

916 <void*>tensor, 

917 DLPACK_VERSIONED_TENSOR_UNUSED_NAME, 

918 _smv_pycapsule_deleter, 

919 ) 

920 py_view = _smv_from_dlpack_capsule(capsule, capsule) 1 ctx1b

921 cpython.Py_INCREF(py_view) 1 ctx1b

922 out_py_object[0] = <void*>py_view 1 ctx1b

923 except Exception: 

924 return -1 

925 return 0 1 ctx1b

926  

927  

928cdef int _smv_dltensor_from_py_object_no_sync( (empty)

929 void* py_object, 

930 DLTensor* out, 

931) noexcept with gil: 

932 if out == NULL: 2 ctx2XbA

933 cpython.PyErr_SetString(RuntimeError, b"out cannot be NULL") 

934 return -1 

935 cdef object obj = <object>py_object 2 ctx2XbA

936 if not isinstance(obj, StridedMemoryView): 2 ctx2XbA

937 cpython.PyErr_SetString(TypeError, b"py_object must be a StridedMemoryView") 1 ctx2Xb

938 return -1 1 ctx2Xb

939 try: 1 ctx1A

940 _smv_setup_dltensor_borrowed(out, <StridedMemoryView>obj) 1 ctx1A

941 except Exception: 

942 return -1 

943 return 0 1 ctx1A

944  

945  

946cdef int _smv_current_work_stream( (empty)

947 _DLDeviceType device_type, 

948 int32_t device_id, 

949 void** out_current_stream, 

950) noexcept with gil: 

951 if out_current_stream == NULL: 1 ctx22b

952 cpython.PyErr_SetString(RuntimeError, b"out_current_stream cannot be NULL") 

953 return -1 

954 # cuda.core has no global/current stream state today. 

955 out_current_stream[0] = NULL 1 ctx22b

956 return 0 1 ctx22b

957  

958  

959cdef void _init_smv_dlpack_exchange_api(): (empty)

960 global _SMV_DLPACK_EXCHANGE_API_INITED 

961 if _SMV_DLPACK_EXCHANGE_API_INITED: (empty)

962 return 

963 _SMV_DLPACK_EXCHANGE_API.header.version.major = DLPACK_MAJOR_VERSION (empty)

964 _SMV_DLPACK_EXCHANGE_API.header.version.minor = DLPACK_MINOR_VERSION (empty)

965 _SMV_DLPACK_EXCHANGE_API.header.prev_api = NULL (empty)

966 _SMV_DLPACK_EXCHANGE_API.managed_tensor_allocator = _smv_managed_tensor_allocator (empty)

967 _SMV_DLPACK_EXCHANGE_API.managed_tensor_from_py_object_no_sync = _smv_managed_tensor_from_py_object_no_sync (empty)

968 _SMV_DLPACK_EXCHANGE_API.managed_tensor_to_py_object_no_sync = _smv_managed_tensor_to_py_object_no_sync (empty)

969 _SMV_DLPACK_EXCHANGE_API.dltensor_from_py_object_no_sync = _smv_dltensor_from_py_object_no_sync (empty)

970 _SMV_DLPACK_EXCHANGE_API.current_work_stream = _smv_current_work_stream (empty)

971 _SMV_DLPACK_EXCHANGE_API_INITED = True (empty)

972  

973  

974_init_smv_dlpack_exchange_api() (empty)

975# cdef classes are immutable types in Cython 3, so inject these attributes 

976# directly into the type dict. 

977(<dict>(<PyTypeObject*>StridedMemoryView).tp_dict)["__dlpack_c_exchange_api__"] = _SMV_DLPACK_EXCHANGE_API_CAPSULE (empty)

978(<dict>(<PyTypeObject*>StridedMemoryView).tp_dict)["__c_dlpack_exchange_api__"] = _SMV_DLPACK_EXCHANGE_API_CAPSULE (empty)

979PyType_Modified(<PyTypeObject*>StridedMemoryView) (empty)

980  

981  

982cdef str get_simple_repr(obj): (empty)

983 # TODO: better handling in np.dtype objects 

984 cdef object obj_class 

985 cdef str obj_repr 

986 if isinstance(obj, type): 2 ctx19I

987 obj_class = obj 

988 else: 

989 obj_class = obj.__class__ 2 ctx19I

990 if obj_class.__module__ in (None, "builtins"): 2 ctx19I

991 obj_repr = obj_class.__name__ 1 ctx19

992 else: 

993 obj_repr = f"{obj_class.__module__}.{obj_class.__name__}" 2 ctx19I

994 return obj_repr 2 ctx19I

995  

996  

997  

998cdef bint check_has_dlpack(obj) except*: (empty)

999 cdef bint has_dlpack 

1000 if hasattr(obj, "__dlpack__") and hasattr(obj, "__dlpack_device__"): 54 ctx21 2 3 8 V P J K L M N Q R S T U W X Y Z 0 Zbx y w 6 7 c } TbI 5 H A b s e f h i j o p q r k l m n g u t d a

1001 has_dlpack = True 51 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HAbsefhijopqrklmngutda

1002 elif hasattr(obj, "__cuda_array_interface__"): 3 ctx2Zb} Tb

1003 has_dlpack = False 2 ctx2} Tb

1004 else: 

1005 raise BufferError( 1 ctx2Zb

1006 "the input object does not support any data exchange protocol") 

1007 return has_dlpack 53 ctx21 2 3 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c } TbI 5 H A b s e f h i j o p q r k l m n g u t d a

1008  

1009  

1010cdef class _StridedMemoryViewProxy: 

1011 cdef readonly: 

1012 object obj 

1013 bint has_dlpack 

1014  

1015 def __init__(self, obj: object) -> None: 

1016 self.obj = obj 6 ctx2J K L M N Tb

1017 self.has_dlpack = check_has_dlpack(obj) 6 ctx2J K L M N Tb

1018  

1019 cpdef StridedMemoryView view(self, stream_ptr=None): (empty)

1020 if self.has_dlpack: 5 ctx1JKLMN

1021 return StridedMemoryView.from_dlpack(self.obj, stream_ptr) 5 ctx1JKLMN

1022 else: 

1023 return StridedMemoryView.from_cuda_array_interface(self.obj, stream_ptr) 

1024  

1025  

1026cdef StridedMemoryView view_as_dlpack(obj, stream_ptr, view=None): (empty)

1027 cdef int dldevice, device_id 

1028 cdef bint is_device_accessible, is_readonly 

1029 is_device_accessible = False 59 ctx21 2 3 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

1030 dldevice, device_id = obj.__dlpack_device__() 59 ctx21 2 3 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

1031 if dldevice == _kDLCPU: 59 ctx21 2 3 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba Mb

1032 assert device_id == 0 56 ctx23 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba

1033 device_id = -1 56 ctx23 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba

1034 if stream_ptr is None: 56 ctx23 8 V P J K L M N Q R S T U W X Y Z 0 x y w 6 7 c I 5 H B C D E F G A b s e f h i j o p q r k l m n g u t d Lba

1035 raise BufferError("stream=None is ambiguous with view()") 1 ctx2Lb

1036 elif stream_ptr == -1: 55 ctx138VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1037 stream_ptr = None 55 ctx138VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1038 elif dldevice == _kDLCUDA: 

1039 assert device_id >= 0 

1040 is_device_accessible = True 

1041 # no need to check other stream values, it's a pass-through 

1042 if stream_ptr is None: 

1043 raise BufferError("stream=None is ambiguous with view()") 

1044 elif dldevice in (_kDLCUDAHost, _kDLCUDAManaged): 

1045 is_device_accessible = True 2 ctx112

1046 # just do a pass-through without any checks, as pinned/managed memory can be 

1047 # accessed on both host and device 

1048 else: 

1049 raise BufferError("device not supported") 1 ctx2Mb

1050  

1051 cdef object capsule 

1052 try: 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1053 capsule = obj.__dlpack__( 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1054 stream=int(stream_ptr) if stream_ptr else None, 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1055 max_version=(DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION)) 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1056 except TypeError: 1 ctx1a

1057 capsule = obj.__dlpack__( 1 ctx1a

1058 stream=int(stream_ptr) if stream_ptr else None) 1 ctx1a

1059  

1060 cdef void* data = NULL 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1061 cdef DLTensor* dl_tensor 

1062 cdef DLManagedTensorVersioned* dlm_tensor_ver 

1063 cdef DLManagedTensor* dlm_tensor 

1064 cdef const char *used_name 

1065 if cpython.PyCapsule_IsValid( 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1066 capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME): 

1067 data = cpython.PyCapsule_GetPointer( 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1068 capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME) 

1069 dlm_tensor_ver = <DLManagedTensorVersioned*>data 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1070 dl_tensor = &dlm_tensor_ver.dl_tensor 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1071 is_readonly = bool((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0) 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1072 used_name = DLPACK_VERSIONED_TENSOR_USED_NAME 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1073 elif cpython.PyCapsule_IsValid( 1 ctx1a

1074 capsule, DLPACK_TENSOR_UNUSED_NAME): 

1075 data = cpython.PyCapsule_GetPointer( 1 ctx1a

1076 capsule, DLPACK_TENSOR_UNUSED_NAME) 

1077 dlm_tensor = <DLManagedTensor*>data 1 ctx1a

1078 dl_tensor = &dlm_tensor.dl_tensor 1 ctx1a

1079 is_readonly = False 1 ctx1a

1080 used_name = DLPACK_TENSOR_USED_NAME 1 ctx1a

1081 else: 

1082 assert False 

1083  

1084 cpython.PyCapsule_SetName(capsule, used_name) 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1085  

1086 cdef StridedMemoryView buf = StridedMemoryView() if view is None else view 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1087 buf.dl_tensor = dl_tensor 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1088 buf.metadata = capsule 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1089 buf.ptr = <intptr_t>(dl_tensor.data) 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1090 buf.device_id = device_id 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1091 buf.is_device_accessible = is_device_accessible 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1092 buf.readonly = is_readonly 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1093 buf.exporting_obj = obj 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1094  

1095 return buf 57 ctx11238VPJKLMNQRSTUWXYZ0xyw67cI5HBCDEFGAbsefhijopqrklmngutda

1096  

1097  

1098@functools.lru_cache (empty)

1099def _typestr2dtype(str typestr) -> numpy.dtype: 

1100 return numpy.dtype(typestr) 11 ctx1!#$%'()*+,}

1101  

1102  

1103@functools.lru_cache (empty)

1104def _typestr2itemsize(str typestr) -> int: 

1105 return _typestr2dtype(typestr).itemsize 11 ctx1!#$%'()*+,}

1106  

1107  

1108cdef object dtype_dlpack_to_numpy(DLDataType* dtype): (empty)

1109 cdef int bits = dtype.bits 29 ctx1cIBCDEFGAbsefhijopqrklmngutda

1110 if dtype.lanes != 1: 29 ctx1cIBCDEFGAbsefhijopqrklmngutda

1111 # TODO: return a NumPy structured dtype? 

1112 raise NotImplementedError( 

1113 f'vector dtypes (lanes={dtype.lanes}) is not supported') 

1114 if dtype.code == kDLUInt: 29 ctx1cIBCDEFGAbsefhijopqrklmngutda

1115 if bits == 8: 4 ctx1klmn

1116 np_dtype = numpy.uint8 1 ctx1k

1117 elif bits == 16: 

1118 np_dtype = numpy.uint16 1 ctx1l

1119 elif bits == 32: