Coverage for cuda/core/_tensor

3# SPDX-License-Identifier: Apache-2.0

5from libc.stdint cimport intptr_t, int64_t, uint32_t, uint64_t

6from libc.stddef cimport size_t

7from cuda.bindings cimport cydriver

8from cuda.core._utils.cuda_utils cimport HANDLE_RETURN

9from cuda.core._dlpack cimport (

10 DLDataType, DLTensor, kDLInt, kDLUInt, kDLFloat, kDLBfloat, _DLDeviceType, _kDLCUDA

11)

12from cuda.core._layout cimport _StridedLayout

13from cuda.core._memoryview cimport StridedMemoryView

15import enum

16from dataclasses import dataclass

17from typing import TYPE_CHECKING

19import numpy

21from cuda.core._utils.cuda_utils import CUDAError, check_or_create_options

23if TYPE_CHECKING:

24 from cuda.core._device import Device

26cdef extern from "<cuda/std/span>" nogil:

27 cdef cppclass const_int_span "::cuda::std::span<const int>":

28 const_int_span(const int*, size_t)

31cdef extern from "<cuda/tma>" namespace "cuda" nogil:

32 # Only the `none` enumerators are spelled out: the options arrive as the

33 # CU_TENSOR_MAP_* values these scoped enums are defined from (see the enums

34 # below), so they are cast rather than mapped name by name.

35 cdef enum class tma_interleave_layout(int):

36 none

37 cdef enum class tma_swizzle(int):

38 none

39 cdef enum class tma_l2_fetch_size(int):

40 none

41 cdef enum class tma_oob_fill(int):

42 none

44 cydriver.CUtensorMap make_tma_descriptor(

45 const DLTensor&,

46 const_int_span box_sizes,

47 const_int_span elem_strides,

48 tma_interleave_layout,

49 tma_swizzle,

50 tma_l2_fetch_size,

51 tma_oob_fill) except +

54try:

55 from ml_dtypes import bfloat16 as ml_bfloat16

56except ImportError:

57 ml_bfloat16 = None

59__all__ = ['TensorMapDescriptor', 'TensorMapDescriptorOptions']

62class TensorMapDataType(enum.IntEnum):

63 """Data types for tensor map descriptors.

65 These correspond to the ``CUtensorMapDataType`` driver enum values.

66 """

67 UINT8 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8

68 UINT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16

69 UINT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32

70 INT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32

71 UINT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64

72 INT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64

73 FLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16

74 FLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32

75 FLOAT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64

76 BFLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16

77 FLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ

78 TFLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32

79 TFLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ

82class TensorMapInterleave(enum.IntEnum):

83 """Interleave layout for tensor map descriptors.

85 These correspond to the ``CUtensorMapInterleave`` driver enum values.

86 """

87 NONE = cydriver.CU_TENSOR_MAP_INTERLEAVE_NONE

88 INTERLEAVE_16B = cydriver.CU_TENSOR_MAP_INTERLEAVE_16B

89 INTERLEAVE_32B = cydriver.CU_TENSOR_MAP_INTERLEAVE_32B

92class TensorMapSwizzle(enum.IntEnum):

93 """Swizzle mode for tensor map descriptors.

95 These correspond to the ``CUtensorMapSwizzle`` driver enum values.

96 """

97 NONE = cydriver.CU_TENSOR_MAP_SWIZZLE_NONE

98 SWIZZLE_32B = cydriver.CU_TENSOR_MAP_SWIZZLE_32B

99 SWIZZLE_64B = cydriver.CU_TENSOR_MAP_SWIZZLE_64B

100 SWIZZLE_128B = cydriver.CU_TENSOR_MAP_SWIZZLE_128B

101

102

103class TensorMapL2Promotion(enum.IntEnum):

104 """L2 promotion mode for tensor map descriptors.

105

106 These correspond to the ``CUtensorMapL2promotion`` driver enum values.

107 """

108 NONE = cydriver.CU_TENSOR_MAP_L2_PROMOTION_NONE

109 L2_64B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_64B

110 L2_128B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_128B

111 L2_256B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_256B

112

113

114class TensorMapOOBFill(enum.IntEnum):

115 """Out-of-bounds fill mode for tensor map descriptors.

116

117 These correspond to the ``CUtensorMapFloatOOBfill`` driver enum values.

118 """

119 NONE = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE

120 NAN_REQUEST_ZERO_FMA = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA

121

122

123IF CUDA_CORE_BUILD_MAJOR >= 13:

124 class TensorMapIm2ColWideMode(enum.IntEnum):

125 """Im2col wide mode for tensor map descriptors.

126

127 These correspond to the ``CUtensorMapIm2ColWideMode`` driver enum values.

128 Supported on compute capability 10.0+.

129 """

130 W = cydriver.CU_TENSOR_MAP_IM2COL_WIDE_MODE_W

131 W128 = cydriver.CU_TENSOR_MAP_IM2COL_WIDE_MODE_W128

132ELSE:

133 class TensorMapIm2ColWideMode(enum.IntEnum):

134 """Im2col wide mode for tensor map descriptors.

135

136 This enum is always defined for API stability, but the

137 :meth:`TensorMapDescriptor._from_im2col_wide` factory requires a CUDA 13+

138 build and will raise otherwise.

139 """

140 W = 0

141 W128 = 1

142

143

144_TMA_DT_UINT8: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8)

145_TMA_DT_UINT16: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16)

146_TMA_DT_UINT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32)

147_TMA_DT_INT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32)

148_TMA_DT_UINT64: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64)

149_TMA_DT_INT64: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64)

150_TMA_DT_FLOAT16: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16)

151_TMA_DT_FLOAT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32)

152_TMA_DT_FLOAT64: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64)

153_TMA_DT_BFLOAT16: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16)

154_TMA_DT_FLOAT32_FTZ: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ)

155_TMA_DT_TFLOAT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32)

156_TMA_DT_TFLOAT32_FTZ: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ)

157

158

159def _normalize_tensor_map_data_type(data_type):

160 if data_type is None or isinstance(data_type, TensorMapDataType): 8 ctx1feabcdhg

161 return data_type 7 ctx1feabdhg

162 try: 1 ctx1c

163 return numpy.dtype(data_type) 1 ctx1c

164 except TypeError as e:

165 raise TypeError(

166 "data_type must be a TensorMapDataType or a numpy/ml_dtypes dtype, "

167 f"got {type(data_type)}") from e

168

169

170def _normalize_tensor_map_sequence(name, values):

171 try: 10 ctx1feabcdijhg

172 values = tuple(values) 10 ctx1feabcdijhg

173 except TypeError as e: 1 ctx1j

174 raise TypeError(f"{name} must be a tuple of ints, got {type(values)}") from e 1 ctx1j

175 for i, value in enumerate(values): 9 ctx1feabcdihg

176 if not isinstance(value, int): 9 ctx1feabcdihg

177 raise TypeError(f"{name}[{i}] must be an int, got {type(value)}") 1 ctx1i

178 return values 8 ctx1feabcdhg

179

180

181def _require_tensor_map_enum(name, value, enum_type):

182 if not isinstance(value, enum_type): 8 ctx1feabcdhg

183 raise TypeError(f"{name} must be a {enum_type.__name__}, got {type(value)}") 2 ctx1hg

184 return value 7 ctx1feabcdg

185

186

187@dataclass

188class TensorMapDescriptorOptions:

189 """Options for :meth:`cuda.core.StridedMemoryView.as_tensor_map`.

190

191 Attributes

192 ----------

193 box_dim : tuple[int, ...]

194 Tile size for each tensor dimension, expressed in elements.

195 element_strides : tuple[int, ...], optional

196 Per-dimension element traversal strides.

197 data_type : object, optional

198 Explicit dtype override. Prefer NumPy or ``ml_dtypes`` dtype objects;

199 :class:`TensorMapDataType` remains accepted for compatibility.

200 interleave : TensorMapInterleave, optional

201 Interleave layout. Default ``NONE``.

202 swizzle : TensorMapSwizzle, optional

203 Swizzle mode. Default ``NONE``.

204 l2_promotion : TensorMapL2Promotion, optional

205 L2 promotion mode. Default ``NONE``.

206 oob_fill : TensorMapOOBFill, optional

207 Out-of-bounds fill mode. Default ``NONE``.

208 """

209

210 box_dim: tuple[int, ...]

211 element_strides: tuple[int, ...] | None = None

212 data_type: object = None

213 interleave: TensorMapInterleave = TensorMapInterleave.NONE

214 swizzle: TensorMapSwizzle = TensorMapSwizzle.NONE

215 l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE

216 oob_fill: TensorMapOOBFill = TensorMapOOBFill.NONE

217

218 def __post_init__(self) -> None:

219 self.box_dim = _normalize_tensor_map_sequence("box_dim", self.box_dim) 10 ctx1feabcdijhg

220 if self.element_strides is not None: 8 ctx1feabcdhg

221 self.element_strides = _normalize_tensor_map_sequence("element_strides", self.element_strides) 3 ctx1fab

222 self.data_type = _normalize_tensor_map_data_type(self.data_type) 8 ctx1feabcdhg

223 self.interleave = _require_tensor_map_enum("interleave", self.interleave, TensorMapInterleave) 8 ctx1feabcdhg

224 self.swizzle = _require_tensor_map_enum("swizzle", self.swizzle, TensorMapSwizzle) 7 ctx1feabcdg

225 self.l2_promotion = _require_tensor_map_enum("l2_promotion", self.l2_promotion, TensorMapL2Promotion) 6 ctx1feabcd

226 self.oob_fill = _require_tensor_map_enum("oob_fill", self.oob_fill, TensorMapOOBFill) 6 ctx1feabcd

227

228

229def _coerce_tensor_map_descriptor_options(

230 box_dim,

231 options,

232 *,

233 element_strides,

234 data_type,

235 interleave,

236 swizzle,

237 l2_promotion,

238 oob_fill,

239):

240 if options is not None: 7 ctx1feabcdk

241 if (

242 box_dim is not None

243 or element_strides is not None

244 or data_type is not None

245 or interleave != TensorMapInterleave.NONE

246 or swizzle != TensorMapSwizzle.NONE

247 or l2_promotion != TensorMapL2Promotion.NONE

248 or oob_fill != TensorMapOOBFill.NONE

249 ):

250 raise TypeError(

251 "Specify either options or the individual tensor map arguments, not both")

252 return check_or_create_options(

253 TensorMapDescriptorOptions,

254 options,

255 "Tensor map descriptor options",

256 )

257

258 if box_dim is None: 7 ctx1feabcdk

259 raise TypeError("box_dim is required unless options is provided") 1 ctx1k

260

261 return TensorMapDescriptorOptions( 6 ctx1feabcd

262 box_dim=box_dim,

263 element_strides=element_strides,

264 data_type=data_type,

265 interleave=interleave,

266 swizzle=swizzle,

267 l2_promotion=l2_promotion,

268 oob_fill=oob_fill, 6 ctx1feabcd

269 )

270

271

272# Mapping from numpy dtype to TMA data type

273_NUMPY_DTYPE_TO_TMA = {

274 numpy.dtype(numpy.uint8): _TMA_DT_UINT8,

275 numpy.dtype(numpy.uint16): _TMA_DT_UINT16,

276 numpy.dtype(numpy.uint32): _TMA_DT_UINT32,

277 numpy.dtype(numpy.int32): _TMA_DT_INT32,

278 numpy.dtype(numpy.uint64): _TMA_DT_UINT64,

279 numpy.dtype(numpy.int64): _TMA_DT_INT64,

280 numpy.dtype(numpy.float16): _TMA_DT_FLOAT16,

281 numpy.dtype(numpy.float32): _TMA_DT_FLOAT32,

282 numpy.dtype(numpy.float64): _TMA_DT_FLOAT64,

283}

284

285if ml_bfloat16 is not None:

286 _NUMPY_DTYPE_TO_TMA[numpy.dtype(ml_bfloat16)] = _TMA_DT_BFLOAT16

287

288

289# Mapping from TMA data type to element size in bytes

290_TMA_DATA_TYPE_SIZE = {

291 _TMA_DT_UINT8: 1,

292 _TMA_DT_UINT16: 2,

293 _TMA_DT_UINT32: 4,

294 _TMA_DT_INT32: 4,

295 _TMA_DT_UINT64: 8,

296 _TMA_DT_INT64: 8,

297 _TMA_DT_FLOAT16: 2,

298 _TMA_DT_FLOAT32: 4,

299 _TMA_DT_FLOAT64: 8,

300 _TMA_DT_BFLOAT16: 2,

301 _TMA_DT_FLOAT32_FTZ: 4,

302 _TMA_DT_TFLOAT32: 4,

303 _TMA_DT_TFLOAT32_FTZ: 4,

304}

305

306cdef _resolve_data_type(StridedMemoryView view, data_type):

307 """Resolve the TMA data type from an explicit value or the view's dtype."""

308

309 if data_type is not None: 5 ctx1eabcd

310 if isinstance(data_type, TensorMapDataType): 3 ctx1abc

311 return int(data_type) 2 ctx1ab

312 dt = _normalize_tensor_map_data_type(data_type) 1 ctx1c

313 tma_dt = _NUMPY_DTYPE_TO_TMA.get(dt) 1 ctx1c

314 if tma_dt is None: 1 ctx1c

315 raise ValueError( 1 ctx1c

316 f"Unsupported dtype {dt} for TMA; " 1 ctx1c

317 f"supported dtypes: {list(_NUMPY_DTYPE_TO_TMA.keys())}.") 1 ctx1c

318 return tma_dt

319

320 dt = view.get_dtype() 2 ctx1ed

321 if dt is None: 2 ctx1ed

322 raise ValueError( 1 ctx1e

323 "Cannot infer TMA data type from the tensor; "

324 "please specify data_type explicitly")

325

326 tma_dt = _NUMPY_DTYPE_TO_TMA.get(dt) 1 ctx1d

327 if tma_dt is None: 1 ctx1d

328 raise ValueError( 1 ctx1d

329 f"Unsupported dtype {dt} for TMA; " 1 ctx1d

330 f"supported dtypes: {list(_NUMPY_DTYPE_TO_TMA.keys())}. " 1 ctx1d

331 "You may also specify data_type explicitly.")

332

333 return tma_dt

334

335

336cdef inline bint _tma_dtype_to_dlpack(int tma_dt, DLDataType* out) noexcept:

337 """Write the DLPack spelling of a TMA data type, or return False if it has none."""

338 out.lanes = 1

339 if tma_dt == _TMA_DT_UINT8:

340 out.code = kDLUInt

341 out.bits = 8

342 elif tma_dt == _TMA_DT_UINT16:

343 out.code = kDLUInt

344 out.bits = 16

345 elif tma_dt == _TMA_DT_UINT32:

346 out.code = kDLUInt

347 out.bits = 32

348 elif tma_dt == _TMA_DT_UINT64:

349 out.code = kDLUInt

350 out.bits = 64

351 elif tma_dt == _TMA_DT_INT32:

352 out.code = kDLInt

353 out.bits = 32

354 elif tma_dt == _TMA_DT_INT64:

355 out.code = kDLInt

356 out.bits = 64

357 elif tma_dt == _TMA_DT_FLOAT16:

358 out.code = kDLFloat

359 out.bits = 16

360 elif tma_dt == _TMA_DT_FLOAT32:

361 out.code = kDLFloat

362 out.bits = 32

363 elif tma_dt == _TMA_DT_FLOAT64:

364 out.code = kDLFloat

365 out.bits = 64

366 elif tma_dt == _TMA_DT_BFLOAT16:

367 out.code = kDLBfloat

368 out.bits = 16

369 else:

370 return False

371 return True

372

373

374cdef inline int _validate_tensor_map_view(StridedMemoryView view) except -1:

375 if not view.is_device_accessible: 6 ctx1feabcd

376 raise ValueError("The tensor must be device-accessible") 1 ctx1f

377

378 if view.ptr % 16 != 0: 5 ctx1eabcd

379 raise ValueError(

380 f"Global memory address must be 16-byte aligned, "

381 f"got address 0x{view.ptr:x}")

382 return 0 5 ctx1eabcd

383

384

385cdef StridedMemoryView _get_validated_view(tensor):

386 """Obtain a device-accessible StridedMemoryView with a 16-byte-aligned pointer."""

387 cdef StridedMemoryView view

388 if isinstance(tensor, StridedMemoryView):

389 view = <StridedMemoryView>tensor

390 else:

391 # stream_ptr=-1: no stream synchronization needed because descriptor

392 # creation only reads tensor metadata, it does not move data.

393 view = StridedMemoryView.from_any_interface(tensor, stream_ptr=-1)

394 _validate_tensor_map_view(view)

395 return view

396

397

398cdef int _require_view_device(

399 StridedMemoryView view, int expected_device_id, operation) except -1:

400 """Ensure device-local tensors match the current CUDA device.

401

402 DLPack reports host/managed CUDA memory as ``kDLCUDAHost`` /

403 ``kDLCUDAManaged`` with ``device_id=0`` regardless of the current device,

404 so only true ``kDLCUDA`` tensors are rejected by device-id mismatch.

405 """

406 device_type, device_id = view.__dlpack_device__() 5 ctx1eabcd

407 if device_type == _kDLCUDA and device_id != expected_device_id: 5 ctx1eabcd

408 raise ValueError(

409 f"{operation} expects tensor on device {expected_device_id}, got {device_id}")

410 return 0 5 ctx1eabcd

411

412

413cdef inline intptr_t _get_current_context_ptr() except? 0:

414 cdef cydriver.CUcontext ctx

415 with nogil: 5 ctx1eabcd

416 HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx)) 5 ctx1eabcd

417 if ctx == NULL: 5 ctx1eabcd

418 raise RuntimeError("TensorMapDescriptor requires an active CUDA context")

419 return <intptr_t>ctx 5 ctx1eabcd

420

421

422cdef inline int _get_current_device_id() except -1:

423 cdef cydriver.CUdevice dev

424 with nogil: 5 ctx1eabcd

425 HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev)) 5 ctx1eabcd

426 return <int>dev 5 ctx1eabcd

427

428cdef void _fill_global_strides(

429 uint64_t* out, const int64_t* shape, const int64_t* strides,

430 int rank, int elem_size) noexcept nogil:

431 """Write cuTensorMap's globalStrides: byte strides in column-major order.

432

433 Only rank - 1 entries are written: the innermost stride is implicitly the

434 element size and is not part of the array. ``strides`` is NULL for a

435 C-contiguous tensor, matching the layout's own convention.

436 """

437 cdef int i

438 cdef int64_t stride

439 if strides != NULL:

440 for i in range(rank - 1):

441 out[i] = <uint64_t>(strides[rank - 2 - i] * elem_size)

442 return

443

444 # C-contiguous: each stride spans everything nested inside that dimension.

445 stride = elem_size

446 for i in range(rank - 1):

447 stride *= shape[rank - 1 - i]

448 out[i] = <uint64_t>stride

449

450

451cdef int _fill_element_strides(int* out, element_strides, int rank) except -1:

452 """Write the element traversal strides in the tensor's own (row-major) order.

453

454 Rejecting non-positive strides here keeps the driver paths' narrowing to

455 ``uint32_t`` lossless; CCCL and the driver enforce the real TMA limits.

456 """

457 cdef int i

458 cdef int stride

459 if element_strides is None: 2 ctx1ab

460 for i in range(rank):

461 out[i] = 1

462 return 0

463

464 if len(element_strides) != rank: 2 ctx1ab

465 raise ValueError(

466 f"element_strides must have {rank} elements, got {len(element_strides)}")

467 for i in range(rank): 2 ctx1ab

468 # Cython raises OverflowError if the value does not fit in a C int.

469 stride = element_strides[i] 2 ctx1ab

470 if stride < 1: 2 ctx1ab

471 raise ValueError(f"element_strides[{i}] must be positive, got {stride}") 2 ctx1ab

472 out[i] = stride

473 return 0

474

475

476cdef class TensorMapDescriptor:

477 """Describes a TMA (Tensor Memory Accelerator) tensor map for Hopper+ GPUs.

478

479 A ``TensorMapDescriptor`` wraps the opaque 128-byte ``CUtensorMap`` struct

480 used by the hardware TMA unit for efficient bulk data movement between

481 global and shared memory.

482

483 Public tiled descriptors are created via

484 :meth:`cuda.core.StridedMemoryView.as_tensor_map`. Specialized

485 ``_from_*`` helpers remain private while this API surface settles, and

486 descriptors can be passed directly to :func:`~cuda.core.launch` as a

487 kernel argument.

488 """

489

490 def __init__(self):

491 raise RuntimeError( 1 ctx1l

492 "TensorMapDescriptor cannot be instantiated directly. "

493 "Use StridedMemoryView.as_tensor_map() instead.")

494

495 cdef void* _get_data_ptr(self):

496 return <void*>&self._tensor_map

497

498 cdef int _check_context_compat(self) except -1:

499 cdef cydriver.CUcontext current_ctx

500 cdef cydriver.CUdevice current_dev

501 if self._context == 0 and self._device_id < 0:

502 return 0

503 with nogil:

504 HANDLE_RETURN(cydriver.cuCtxGetCurrent(&current_ctx))

505 if current_ctx == NULL:

506 raise RuntimeError("TensorMapDescriptor requires an active CUDA context")

507 if self._context != 0 and <intptr_t>current_ctx != self._context:

508 raise RuntimeError(

509 "TensorMapDescriptor was created in a different CUDA context")

510 with nogil:

511 HANDLE_RETURN(cydriver.cuCtxGetDevice(&current_dev))

512 cdef int current_dev_id = <int>current_dev

513 if self._device_id >= 0 and current_dev_id != self._device_id:

514 raise RuntimeError(

515 f"TensorMapDescriptor belongs to device {self._device_id}, "

516 f"but current device is {current_dev_id}")

517 return 0

518

519 @property

520 def device(self) -> Device | None:

521 """Return the :obj:`~cuda.core.Device` associated with this descriptor."""

522 if self._device_id >= 0:

523 from cuda.core._device import Device

524 return Device(self._device_id)

525 return None

526

527 @classmethod

528 def _from_tiled(cls, StridedMemoryView view, box_dim=None, *,

529 options=None,

530 element_strides=None,

531 data_type=None,

532 interleave=TensorMapInterleave.NONE,

533 swizzle=TensorMapSwizzle.NONE,

534 l2_promotion=TensorMapL2Promotion.NONE,

535 oob_fill=TensorMapOOBFill.NONE):

536 """Create a tiled TMA descriptor from a validated view.

537

538 Parameters

539 ----------

540 view : StridedMemoryView

541 A device-accessible view with a 16-byte-aligned pointer.

542 box_dim : tuple of int, optional

543 The size of each tile dimension (in elements). Must have the

544 same rank as the tensor and each value must be in [1, 256].

545 Specified in the same (row-major) order as the tensor shape.

546 Required unless ``options`` is provided.

547 options : TensorMapDescriptorOptions or mapping, optional

548 Bundled tiled-descriptor options. When provided, do not also pass

549 ``box_dim`` or the individual option kwargs.

550 element_strides : tuple of int, optional

551 Per-dimension element traversal strides. Default is all 1s.

552 Specified in the same (row-major) order as the tensor shape.

553 data_type : dtype-like or TensorMapDataType, optional

554 Explicit dtype override. If ``None``, inferred from the tensor's

555 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is

556 accepted for compatibility.

557 interleave : TensorMapInterleave

558 Interleave layout. Default ``NONE``.

559 swizzle : TensorMapSwizzle

560 Swizzle mode. Default ``NONE``.

561 l2_promotion : TensorMapL2Promotion

562 L2 promotion mode. Default ``NONE``.

563 oob_fill : TensorMapOOBFill

564 Out-of-bounds fill mode. Default ``NONE``.

565

566 Returns

567 -------

568 TensorMapDescriptor

569

570 Raises

571 ------

572 ValueError

573 If the tensor rank is outside [1, 5], the pointer is not

574 16-byte aligned, or dimension/stride constraints are violated.

575 CUDAError

576 If the CUDA driver rejects the encoded descriptor.

577 """

578 cdef TensorMapDescriptor desc = cls.__new__(cls) 6 ctx1feabcd

579

580 opts = _coerce_tensor_map_descriptor_options( 6 ctx1feabcd

581 box_dim,

582 options,

583 element_strides=element_strides,

584 data_type=data_type,

585 interleave=interleave,

586 swizzle=swizzle,

587 l2_promotion=l2_promotion,

588 oob_fill=oob_fill, 6 ctx1feabcd

589 )

590 box_dim = opts.box_dim 6 ctx1feabcd

591 element_strides = opts.element_strides 6 ctx1feabcd

592 data_type = opts.data_type 6 ctx1feabcd

593 interleave = opts.interleave 6 ctx1feabcd

594 swizzle = opts.swizzle 6 ctx1feabcd

595 l2_promotion = opts.l2_promotion 6 ctx1feabcd

596 oob_fill = opts.oob_fill 6 ctx1feabcd

597

598 # Convert options to driver enum:

599 cdef cydriver.CUtensorMapInterleave c_interleave = interleave 6 ctx1feabcd

600 cdef cydriver.CUtensorMapSwizzle c_swizzle = swizzle 6 ctx1feabcd

601 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = l2_promotion 6 ctx1feabcd

602 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = oob_fill 6 ctx1feabcd

603

604 _validate_tensor_map_view(view) 6 ctx1feabcd

605 # Keep both the original tensor object and the validated view alive.

606 # For DLPack exporters, the view may hold the owning capsule whose

607 # deleter can free the backing allocation when released.

608 desc._source_ref = view.exporting_obj 5 ctx1eabcd

609 desc._view_ref = view 5 ctx1eabcd

610 desc._context = _get_current_context_ptr() 5 ctx1eabcd

611 desc._device_id = _get_current_device_id() 5 ctx1eabcd

612 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_tiled") 5 ctx1eabcd

613

614 tma_dt = _resolve_data_type(view, data_type) 5 ctx1eabcd

615 cdef int c_data_type_int = tma_dt 2 ctx1ab

616 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int 2 ctx1ab

617

618 cdef intptr_t global_address = view.ptr 2 ctx1ab

619 # The layout owns the shape/stride arrays in exactly the form DLPack and

620 # the driver want them, so they are read in place instead of going

621 # through the tuple-building Python properties. strides is NULL when the

622 # tensor is C-contiguous.

623 cdef _StridedLayout layout = view.get_layout() 2 ctx1ab

624 cdef const int64_t* view_shape = layout.base.shape 2 ctx1ab

625 cdef const int64_t* view_strides = layout.base.strides 2 ctx1ab

626

627 cdef int rank = layout.base.ndim 2 ctx1ab

628 if rank < 1 or rank > 5: 2 ctx1ab

629 raise ValueError(

630 f"Tensor rank must be between 1 and 5, got {rank}")

631

632 if len(box_dim) != rank: 2 ctx1ab

633 raise ValueError(

634 f"box_dim must have {rank} elements (same as tensor rank), "

635 f"got {len(box_dim)}")

636

637 for i, bd in enumerate(box_dim): 2 ctx1ab

638 if bd < 1 or bd > 256: 2 ctx1ab

639 raise ValueError(

640 f"box_dim[{i}] must be in [1, 256], got {bd}")

641

642 cdef int c_elem_strides[5]

643 _fill_element_strides(c_elem_strides, element_strides, rank) 2 ctx1ab

644

645 # Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible.

646 # This avoids maintaining a second, independent validation/encoding implementation.

647 cdef DLTensor dl_tensor

648 cdef int64_t c_strides[5]

649 cdef int64_t c_stride

650 cdef int c_box_sizes[5]

651 cdef int i_dl

652 cdef int dl_device_type

653 if _tma_dtype_to_dlpack(tma_dt, &dl_tensor.dtype):

654 for i_dl in range(rank):

655 c_box_sizes[i_dl] = <int>box_dim[i_dl]

656

657 if view_strides != NULL:

658 dl_tensor.strides = <int64_t*>view_strides

659 else:

660 # DLPack 1.2 and later reject strides=NULL.

661 c_stride = 1

662 for i_dl in range(rank - 1, -1, -1):

663 c_strides[i_dl] = c_stride

664 c_stride *= view_shape[i_dl]

665 dl_tensor.strides = &c_strides[0]

666

667 dl_device_type = view.__dlpack_device__()[0]

668 dl_tensor.data = <void*>global_address

669 dl_tensor.device.device_type = <_DLDeviceType>dl_device_type

670 # DLPack reports kDLCUDAManaged / kDLCUDAHost as device_id=0.

671 # CCCL uses this id for compute-capability and shared-memory

672 # queries, so pass the current CUDA device instead. kDLCUDA

673 # tensors were already required to match this id.

674 dl_tensor.device.device_id = desc._device_id

675 dl_tensor.ndim = rank

676 dl_tensor.shape = <int64_t*>view_shape

677 dl_tensor.byte_offset = 0

678

679 try:

680 with nogil:

681 desc._tensor_map = make_tma_descriptor(

682 dl_tensor,

683 const_int_span(&c_box_sizes[0], <size_t>rank),

684 const_int_span(&c_elem_strides[0], <size_t>rank),

685 <tma_interleave_layout>c_interleave,

686 <tma_swizzle>c_swizzle,

687 <tma_l2_fetch_size>c_l2_promotion,

688 <tma_oob_fill>c_oob_fill,

689 )

690 except RuntimeError as err:

691 # RuntimeErrors here should in practice be CUDAErrors.

692 # (either driver error or compute capability check).

693 raise CUDAError(str(err)) from err

694 desc._repr_info = {

695 "method": "tiled",

696 "rank": rank,

697 "data_type": TensorMapDataType(tma_dt),

698 "swizzle": swizzle,

699 }

700 return desc

701

702 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]

703

704 # Reverse dimensions for column-major cuTensorMap convention

705 # Python/DLPack: row-major (dim 0 = outermost)

706 # cuTensorMap: column-major (dim 0 = innermost)

707 cdef uint64_t[5] c_global_dim

708 cdef uint64_t[4] c_global_strides # rank - 1 elements

709 cdef uint32_t[5] c_box_dim

710 cdef uint32_t[5] c_element_strides

711 cdef int i_c

712

713 for i_c in range(rank):

714 # Reverse: Python dim i -> cuTensorMap dim (rank - 1 - i)

715 c_global_dim[i_c] = <uint64_t>view_shape[rank - 1 - i_c]

716 c_box_dim[i_c] = <uint32_t>box_dim[rank - 1 - i_c]

717 c_element_strides[i_c] = <uint32_t>c_elem_strides[rank - 1 - i_c]

718

719 _fill_global_strides(c_global_strides, view_shape, view_strides, rank, elem_size)

720

721 cdef uint32_t c_rank = <uint32_t>rank

722

723 with nogil:

724 HANDLE_RETURN(cydriver.cuTensorMapEncodeTiled(

725 &desc._tensor_map,

726 c_data_type,

727 c_rank,

728 <void*>global_address,

729 c_global_dim,

730 c_global_strides,

731 c_box_dim,

732 c_element_strides,

733 c_interleave,

734 c_swizzle,

735 c_l2_promotion,

736 c_oob_fill,

737 ))

738

739 desc._repr_info = {

740 "method": "tiled",

741 "rank": rank,

742 "data_type": TensorMapDataType(tma_dt),

743 "swizzle": swizzle,

744 }

745

746 return desc

747

748 @classmethod

749 def _from_im2col(cls, StridedMemoryView view, pixel_box_lower_corner, pixel_box_upper_corner,

750 channels_per_pixel, pixels_per_column, *,

751 element_strides=None,

752 data_type=None,

753 interleave=TensorMapInterleave.NONE,

754 swizzle=TensorMapSwizzle.NONE,

755 l2_promotion=TensorMapL2Promotion.NONE,

756 oob_fill=TensorMapOOBFill.NONE):

757 """Create an im2col TMA descriptor from a validated view.

758

759 Im2col layout is used for convolution-style data access patterns.

760

761 Parameters

762 ----------

763 view : StridedMemoryView

764 A device-accessible view with a 16-byte-aligned pointer.

765 pixel_box_lower_corner : tuple of int

766 Lower corner of the pixel bounding box for each spatial

767 dimension (rank - 2 elements). Specified in row-major order

768 matching the tensor's spatial dimensions.

769 pixel_box_upper_corner : tuple of int

770 Upper corner of the pixel bounding box for each spatial

771 dimension (rank - 2 elements). Specified in row-major order

772 matching the tensor's spatial dimensions.

773 channels_per_pixel : int

774 Number of channels per pixel.

775 pixels_per_column : int

776 Number of pixels per column.

777 element_strides : tuple of int, optional

778 Per-dimension element traversal strides. Default is all 1s.

779 data_type : dtype-like or TensorMapDataType, optional

780 Explicit dtype override. If ``None``, inferred from the tensor's

781 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is

782 accepted for compatibility.

783 interleave : TensorMapInterleave

784 Interleave layout. Default ``NONE``.

785 swizzle : TensorMapSwizzle

786 Swizzle mode. Default ``NONE``.

787 l2_promotion : TensorMapL2Promotion

788 L2 promotion mode. Default ``NONE``.

789 oob_fill : TensorMapOOBFill

790 Out-of-bounds fill mode. Default ``NONE``.

791

792 Returns

793 -------

794 TensorMapDescriptor

795

796 Raises

797 ------

798 ValueError

799 If the tensor rank is outside [3, 5], the pointer is not

800 16-byte aligned, or other constraints are violated.

801 """

802 cdef TensorMapDescriptor desc = cls.__new__(cls)

803

804 _validate_tensor_map_view(view)

805 desc._source_ref = view.exporting_obj

806 desc._view_ref = view

807 desc._context = _get_current_context_ptr()

808 desc._device_id = _get_current_device_id()

809 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_im2col")

810

811 tma_dt = _resolve_data_type(view, data_type)

812 cdef int c_data_type_int = tma_dt

813 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int

814

815 cdef intptr_t global_address = view.ptr

816 cdef _StridedLayout layout = view.get_layout()

817 cdef const int64_t* view_shape = layout.base.shape

818 cdef const int64_t* view_strides = layout.base.strides

819

820 cdef int rank = layout.base.ndim

821 if rank < 3 or rank > 5:

822 raise ValueError(

823 f"Im2col tensor rank must be between 3 and 5, got {rank}")

824

825 cdef int n_spatial = rank - 2

826 if len(pixel_box_lower_corner) != n_spatial:

827 raise ValueError(

828 f"pixel_box_lower_corner must have {n_spatial} elements "

829 f"(rank - 2), got {len(pixel_box_lower_corner)}")

830 if len(pixel_box_upper_corner) != n_spatial:

831 raise ValueError(

832 f"pixel_box_upper_corner must have {n_spatial} elements "

833 f"(rank - 2), got {len(pixel_box_upper_corner)}")

834

835 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]

836 cdef int c_elem_strides[5]

837 _fill_element_strides(c_elem_strides, element_strides, rank)

838

839 # Reverse all dimension arrays for column-major convention

840 cdef uint64_t[5] c_global_dim

841 cdef uint64_t[4] c_global_strides

842 cdef uint32_t[5] c_element_strides

843 cdef int[3] c_pixel_box_lower # max 3 spatial dims (rank 5 - 2)

844 cdef int[3] c_pixel_box_upper

845 cdef int i_c

846

847 for i_c in range(3):

848 c_pixel_box_lower[i_c] = 0

849 c_pixel_box_upper[i_c] = 0

850

851 for i_c in range(rank):

852 c_global_dim[i_c] = <uint64_t>view_shape[rank - 1 - i_c]

853 c_element_strides[i_c] = <uint32_t>c_elem_strides[rank - 1 - i_c]

854

855 _fill_global_strides(c_global_strides, view_shape, view_strides, rank, elem_size)

856

857 # Reverse spatial dimensions for lower/upper corners

858 for i_c in range(n_spatial):

859 c_pixel_box_lower[i_c] = <int>pixel_box_lower_corner[n_spatial - 1 - i_c]

860 c_pixel_box_upper[i_c] = <int>pixel_box_upper_corner[n_spatial - 1 - i_c]

861

862 cdef uint32_t c_rank = <uint32_t>rank

863 cdef uint32_t c_channels = <uint32_t>channels_per_pixel

864 cdef uint32_t c_pixels = <uint32_t>pixels_per_column

865 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave><int>interleave

866 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle><int>swizzle

867 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion><int>l2_promotion

868 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill><int>oob_fill

869

870 with nogil:

871 HANDLE_RETURN(cydriver.cuTensorMapEncodeIm2col(

872 &desc._tensor_map,

873 c_data_type,

874 c_rank,

875 <void*>global_address,

876 c_global_dim,

877 c_global_strides,

878 c_pixel_box_lower,

879 c_pixel_box_upper,

880 c_channels,

881 c_pixels,

882 c_element_strides,

883 c_interleave,

884 c_swizzle,

885 c_l2_promotion,

886 c_oob_fill,

887 ))

888

889 desc._repr_info = {

890 "method": "im2col",

891 "rank": rank,

892 "data_type": TensorMapDataType(tma_dt),

893 "swizzle": swizzle,

894 }

895

896 return desc

897

898 @classmethod

899 def _from_im2col_wide(cls, StridedMemoryView view,

900 pixel_box_lower_corner_width, pixel_box_upper_corner_width,

901 channels_per_pixel, pixels_per_column, *,

902 element_strides=None,

903 data_type=None,

904 interleave=TensorMapInterleave.NONE,

905 mode=TensorMapIm2ColWideMode.W,

906 swizzle=TensorMapSwizzle.SWIZZLE_128B,

907 l2_promotion=TensorMapL2Promotion.NONE,

908 oob_fill=TensorMapOOBFill.NONE):

909 """Create an im2col-wide TMA descriptor from a validated view.

910

911 Im2col-wide layout loads elements exclusively along the W (width)

912 dimension. This variant is supported on compute capability 10.0+

913 (Blackwell and later).

914

915 Parameters

916 ----------

917 view : StridedMemoryView

918 A device-accessible view with a 16-byte-aligned pointer.

919 pixel_box_lower_corner_width : int

920 Lower corner of the pixel bounding box along the W dimension.

921 pixel_box_upper_corner_width : int

922 Upper corner of the pixel bounding box along the W dimension.

923 channels_per_pixel : int

924 Number of channels per pixel.

925 pixels_per_column : int

926 Number of pixels per column.

927 element_strides : tuple of int, optional

928 Per-dimension element traversal strides. Default is all 1s.

929 data_type : dtype-like or TensorMapDataType, optional

930 Explicit dtype override. If ``None``, inferred from the tensor's

931 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is

932 accepted for compatibility.

933 interleave : TensorMapInterleave

934 Interleave layout. Default ``NONE``.

935 mode : TensorMapIm2ColWideMode

936 Im2col wide mode. Default ``W``.

937 swizzle : TensorMapSwizzle

938 Swizzle mode. Default ``SWIZZLE_128B``.

939 l2_promotion : TensorMapL2Promotion

940 L2 promotion mode. Default ``NONE``.

941 oob_fill : TensorMapOOBFill

942 Out-of-bounds fill mode. Default ``NONE``.

943

944 Returns

945 -------

946 TensorMapDescriptor

947

948 Raises

949 ------

950 ValueError

951 If the tensor rank is outside [3, 5], the pointer is not

952 16-byte aligned, or other constraints are violated.

953 """

954 IF CUDA_CORE_BUILD_MAJOR < 13:

955 raise RuntimeError(

956 "TensorMapDescriptor._from_im2col_wide requires a CUDA 13+ build")

957 ELSE:

958 cdef TensorMapDescriptor desc = cls.__new__(cls)

959

960 _validate_tensor_map_view(view)

961 desc._source_ref = view.exporting_obj

962 desc._view_ref = view

963 desc._context = _get_current_context_ptr()

964 desc._device_id = _get_current_device_id()

965 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_im2col_wide")

966

967 tma_dt = _resolve_data_type(view, data_type)

968 cdef int c_data_type_int = tma_dt

969 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int

970

971 cdef intptr_t global_address = view.ptr

972 cdef _StridedLayout layout = view.get_layout()

973 cdef const int64_t* view_shape = layout.base.shape

974 cdef const int64_t* view_strides = layout.base.strides

975

976 cdef int rank = layout.base.ndim

977 if rank < 3 or rank > 5:

978 raise ValueError(

979 f"Im2col-wide tensor rank must be between 3 and 5, got {rank}")

980

981 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]

982 cdef int c_elem_strides[5]

983 _fill_element_strides(c_elem_strides, element_strides, rank)

984

985 # Reverse all dimension arrays for column-major convention

986 cdef uint64_t[5] c_global_dim

987 cdef uint64_t[4] c_global_strides

988 cdef uint32_t[5] c_element_strides

989 cdef int i_c

990

991 for i_c in range(rank):

992 c_global_dim[i_c] = <uint64_t>view_shape[rank - 1 - i_c]

993 c_element_strides[i_c] = <uint32_t>c_elem_strides[rank - 1 - i_c]

994

995 _fill_global_strides(c_global_strides, view_shape, view_strides, rank, elem_size)

996

997 cdef uint32_t c_rank = <uint32_t>rank

998 cdef int c_lower_w = <int>pixel_box_lower_corner_width

999 cdef int c_upper_w = <int>pixel_box_upper_corner_width

1000 cdef uint32_t c_channels = <uint32_t>channels_per_pixel

1001 cdef uint32_t c_pixels = <uint32_t>pixels_per_column

1002 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave><int>interleave

1003 cdef cydriver.CUtensorMapIm2ColWideMode c_mode = <cydriver.CUtensorMapIm2ColWideMode><int>mode

1004 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle><int>swizzle

1005 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion><int>l2_promotion

1006 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill><int>oob_fill

1007

1008 with nogil:

1009 HANDLE_RETURN(cydriver.cuTensorMapEncodeIm2colWide(

1010 &desc._tensor_map,

1011 c_data_type,

1012 c_rank,

1013 <void*>global_address,

1014 c_global_dim,

1015 c_global_strides,

1016 c_lower_w,

1017 c_upper_w,

1018 c_channels,

1019 c_pixels,

1020 c_element_strides,

1021 c_interleave,

1022 c_mode,

1023 c_swizzle,

1024 c_l2_promotion,

1025 c_oob_fill,

1026 ))

1027

1028 desc._repr_info = {

1029 "method": "im2col_wide",

1030 "rank": rank,

1031 "data_type": TensorMapDataType(tma_dt),

1032 "swizzle": swizzle,

1033 }

1034

1035 return desc

1036

1037 def replace_address(self, tensor: object) -> None:

1038 """Replace the global memory address in this tensor map descriptor.

1039

1040 This is useful when the tensor data has been reallocated but the

1041 shape, strides, and other parameters remain the same.

1042

1043 Parameters

1044 ----------

1045 tensor : object

1046 Any object supporting DLPack or ``__cuda_array_interface__``,

1047 or a :obj:`~cuda.core.StridedMemoryView`. Must refer to

1048 device-accessible memory with a 16-byte-aligned pointer.

1049 """

1050 self._check_context_compat()

1051 cdef StridedMemoryView view = _get_validated_view(tensor)

1052 _require_view_device(view, self._device_id, "replace_address")

1053

1054 cdef intptr_t global_address = view.ptr

1055

1056 with nogil:

1057 HANDLE_RETURN(cydriver.cuTensorMapReplaceAddress(

1058 &self._tensor_map,

1059 <void*>global_address,

1060 ))

1061

1062 # Update the source reference only after the driver call succeeds,

1063 # so we don't drop the old tensor (risking a dangling pointer in the

1064 # CUtensorMap struct) if the call fails.

1065 self._source_ref = view.exporting_obj

1066 self._view_ref = view

1067

1068 def __repr__(self) -> str:

1069 info = self._repr_info

1070 if info is None:

1071 return "TensorMapDescriptor()"

1072 parts = []

1073 if "method" in info:

1074 parts.append(info["method"])

1075 if "rank" in info:

1076 parts.append(f"rank={info['rank']}")

1077 if "data_type" in info:

1078 parts.append(f"dtype={info['data_type'].name}")

1079 if "swizzle" in info:

1080 parts.append(f"swizzle={info['swizzle'].name}")

1081 return f"TensorMapDescriptor({', '.join(parts)})"

Coverage for cuda/core/_tensor_map.pyx: 20.00%

495 statements