Coverage for cuda/core/_tensor_map.pyx: 35.86%
527 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-29 01:38 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from libc.stdint cimport intptr_t, int64_t, uint8_t, uint16_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 kDLInt, kDLUInt, kDLFloat, kDLBfloat, _kDLCUDA
11import enum
12from dataclasses import dataclass
13from typing import TYPE_CHECKING
15import numpy
17from cuda.core._memoryview import StridedMemoryView
18from cuda.core._utils.cuda_utils import check_or_create_options
20if TYPE_CHECKING:
21 from cuda.core._device import Device
23cdef extern from "_cpp/tensor_map_cccl.h":
24 int cuda_core_cccl_make_tma_descriptor_tiled(
25 void* out_tensor_map,
26 void* data,
27 int device_type,
28 int device_id,
29 int ndim,
30 const int64_t* shape,
31 const int64_t* strides,
32 uint8_t dtype_code,
33 uint8_t dtype_bits,
34 uint16_t dtype_lanes,
35 const int* box_sizes,
36 const int* elem_strides,
37 int interleave_layout,
38 int swizzle,
39 int l2_fetch_size,
40 int oob_fill,
41 char* err,
42 size_t err_cap) nogil
45try:
46 from ml_dtypes import bfloat16 as ml_bfloat16
47except ImportError:
48 ml_bfloat16 = None
51class TensorMapDataType(enum.IntEnum):
52 """Data types for tensor map descriptors.
54 These correspond to the ``CUtensorMapDataType`` driver enum values.
55 """
56 UINT8 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8
57 UINT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16
58 UINT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32
59 INT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32
60 UINT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64
61 INT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64
62 FLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16
63 FLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32
64 FLOAT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64
65 BFLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16
66 FLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ
67 TFLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32
68 TFLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ
71class TensorMapInterleave(enum.IntEnum):
72 """Interleave layout for tensor map descriptors.
74 These correspond to the ``CUtensorMapInterleave`` driver enum values.
75 """
76 NONE = cydriver.CU_TENSOR_MAP_INTERLEAVE_NONE
77 INTERLEAVE_16B = cydriver.CU_TENSOR_MAP_INTERLEAVE_16B
78 INTERLEAVE_32B = cydriver.CU_TENSOR_MAP_INTERLEAVE_32B
81class TensorMapSwizzle(enum.IntEnum):
82 """Swizzle mode for tensor map descriptors.
84 These correspond to the ``CUtensorMapSwizzle`` driver enum values.
85 """
86 NONE = cydriver.CU_TENSOR_MAP_SWIZZLE_NONE
87 SWIZZLE_32B = cydriver.CU_TENSOR_MAP_SWIZZLE_32B
88 SWIZZLE_64B = cydriver.CU_TENSOR_MAP_SWIZZLE_64B
89 SWIZZLE_128B = cydriver.CU_TENSOR_MAP_SWIZZLE_128B
92class TensorMapL2Promotion(enum.IntEnum):
93 """L2 promotion mode for tensor map descriptors.
95 These correspond to the ``CUtensorMapL2promotion`` driver enum values.
96 """
97 NONE = cydriver.CU_TENSOR_MAP_L2_PROMOTION_NONE
98 L2_64B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_64B
99 L2_128B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_128B
100 L2_256B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_256B
103class TensorMapOOBFill(enum.IntEnum):
104 """Out-of-bounds fill mode for tensor map descriptors.
106 These correspond to the ``CUtensorMapFloatOOBfill`` driver enum values.
107 """
108 NONE = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
109 NAN_REQUEST_ZERO_FMA = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA
112IF CUDA_CORE_BUILD_MAJOR >= 13:
113 class TensorMapIm2ColWideMode(enum.IntEnum):
114 """Im2col wide mode for tensor map descriptors.
116 These correspond to the ``CUtensorMapIm2ColWideMode`` driver enum values.
117 Supported on compute capability 10.0+.
118 """
119 W = cydriver.CU_TENSOR_MAP_IM2COL_WIDE_MODE_W
120 W128 = cydriver.CU_TENSOR_MAP_IM2COL_WIDE_MODE_W128
121ELSE:
122 class TensorMapIm2ColWideMode(enum.IntEnum):
123 """Im2col wide mode for tensor map descriptors.
125 This enum is always defined for API stability, but the
126 :meth:`TensorMapDescriptor._from_im2col_wide` factory requires a CUDA 13+
127 build and will raise otherwise.
128 """
129 W = 0
130 W128 = 1
133_TMA_DT_UINT8 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8)
134_TMA_DT_UINT16 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16)
135_TMA_DT_UINT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32)
136_TMA_DT_INT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32)
137_TMA_DT_UINT64 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64)
138_TMA_DT_INT64 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64)
139_TMA_DT_FLOAT16 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16)
140_TMA_DT_FLOAT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32)
141_TMA_DT_FLOAT64 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64)
142_TMA_DT_BFLOAT16 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16)
143_TMA_DT_FLOAT32_FTZ = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ)
144_TMA_DT_TFLOAT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32)
145_TMA_DT_TFLOAT32_FTZ = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ)
148def _normalize_tensor_map_data_type(data_type):
149 if data_type is None or isinstance(data_type, TensorMapDataType): 1ab
150 return data_type 1b
151 try:
152 return numpy.dtype(data_type)
153 except TypeError as e:
154 raise TypeError(
155 "data_type must be a TensorMapDataType or a numpy/ml_dtypes dtype, "
156 f"got {type(data_type)}") from e
159def _normalize_tensor_map_sequence(name, values):
160 try: 1b
161 values = tuple(values) 1b
162 except TypeError as e:
163 raise TypeError(f"{name} must be a tuple of ints, got {type(values)}") from e
164 for i, value in enumerate(values): 1ab
165 if not isinstance(value, int): 1b
166 raise TypeError(f"{name}[{i}] must be an int, got {type(value)}")
167 return values 1b
170def _require_tensor_map_enum(name, value, enum_type):
171 if not isinstance(value, enum_type): 1b
172 raise TypeError(f"{name} must be a {enum_type.__name__}, got {type(value)}")
173 return value 1b
176@dataclass
177class TensorMapDescriptorOptions:
178 """Options for :meth:`cuda.core.StridedMemoryView.as_tensor_map`.
180 Attributes
181 ----------
182 box_dim : tuple[int, ...]
183 Tile size for each tensor dimension, expressed in elements.
184 element_strides : tuple[int, ...], optional
185 Per-dimension element traversal strides.
186 data_type : object, optional
187 Explicit dtype override. Prefer NumPy or ``ml_dtypes`` dtype objects;
188 :class:`TensorMapDataType` remains accepted for compatibility.
189 interleave : TensorMapInterleave, optional
190 Interleave layout. Default ``NONE``.
191 swizzle : TensorMapSwizzle, optional
192 Swizzle mode. Default ``NONE``.
193 l2_promotion : TensorMapL2Promotion, optional
194 L2 promotion mode. Default ``NONE``.
195 oob_fill : TensorMapOOBFill, optional
196 Out-of-bounds fill mode. Default ``NONE``.
197 """
199 box_dim: tuple[int, ...]
200 element_strides: tuple[int, ...] | None = None
201 data_type: object = None
202 interleave: TensorMapInterleave = TensorMapInterleave.NONE
203 swizzle: TensorMapSwizzle = TensorMapSwizzle.NONE
204 l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE
205 oob_fill: TensorMapOOBFill = TensorMapOOBFill.NONE
207 def __post_init__(self) -> None:
208 self.box_dim = _normalize_tensor_map_sequence("box_dim", self.box_dim) 1b
209 if self.element_strides is not None: 1b
210 self.element_strides = _normalize_tensor_map_sequence("element_strides", self.element_strides) 1b
211 self.data_type = _normalize_tensor_map_data_type(self.data_type) 1b
212 self.interleave = _require_tensor_map_enum("interleave", self.interleave, TensorMapInterleave) 1b
213 self.swizzle = _require_tensor_map_enum("swizzle", self.swizzle, TensorMapSwizzle) 1ab
214 self.l2_promotion = _require_tensor_map_enum("l2_promotion", self.l2_promotion, TensorMapL2Promotion) 1b
215 self.oob_fill = _require_tensor_map_enum("oob_fill", self.oob_fill, TensorMapOOBFill) 1b
218def _coerce_tensor_map_descriptor_options(
219 box_dim,
220 options,
221 *,
222 element_strides,
223 data_type,
224 interleave,
225 swizzle,
226 l2_promotion,
227 oob_fill,
228):
229 if options is not None: 1b
230 if (
231 box_dim is not None
232 or element_strides is not None
233 or data_type is not None
234 or interleave != TensorMapInterleave.NONE
235 or swizzle != TensorMapSwizzle.NONE
236 or l2_promotion != TensorMapL2Promotion.NONE
237 or oob_fill != TensorMapOOBFill.NONE
238 ):
239 raise TypeError(
240 "Specify either options or the individual tensor map arguments, not both")
241 return check_or_create_options(
242 TensorMapDescriptorOptions,
243 options,
244 "Tensor map descriptor options",
245 )
247 if box_dim is None: 1b
248 raise TypeError("box_dim is required unless options is provided")
250 return TensorMapDescriptorOptions( 1b
251 box_dim=box_dim,
252 element_strides=element_strides,
253 data_type=data_type,
254 interleave=interleave,
255 swizzle=swizzle,
256 l2_promotion=l2_promotion,
257 oob_fill=oob_fill, 1b
258 )
261# Mapping from numpy dtype to TMA data type
262_NUMPY_DTYPE_TO_TMA = {
263 numpy.dtype(numpy.uint8): _TMA_DT_UINT8,
264 numpy.dtype(numpy.uint16): _TMA_DT_UINT16,
265 numpy.dtype(numpy.uint32): _TMA_DT_UINT32,
266 numpy.dtype(numpy.int32): _TMA_DT_INT32,
267 numpy.dtype(numpy.uint64): _TMA_DT_UINT64,
268 numpy.dtype(numpy.int64): _TMA_DT_INT64,
269 numpy.dtype(numpy.float16): _TMA_DT_FLOAT16,
270 numpy.dtype(numpy.float32): _TMA_DT_FLOAT32,
271 numpy.dtype(numpy.float64): _TMA_DT_FLOAT64,
272}
274if ml_bfloat16 is not None:
275 _NUMPY_DTYPE_TO_TMA[numpy.dtype(ml_bfloat16)] = _TMA_DT_BFLOAT16
278# Mapping from TMA data type to element size in bytes
279_TMA_DATA_TYPE_SIZE = {
280 _TMA_DT_UINT8: 1,
281 _TMA_DT_UINT16: 2,
282 _TMA_DT_UINT32: 4,
283 _TMA_DT_INT32: 4,
284 _TMA_DT_UINT64: 8,
285 _TMA_DT_INT64: 8,
286 _TMA_DT_FLOAT16: 2,
287 _TMA_DT_FLOAT32: 4,
288 _TMA_DT_FLOAT64: 8,
289 _TMA_DT_BFLOAT16: 2,
290 _TMA_DT_FLOAT32_FTZ: 4,
291 _TMA_DT_TFLOAT32: 4,
292 _TMA_DT_TFLOAT32_FTZ: 4,
293}
296def _resolve_data_type(view, data_type):
297 """Resolve the TMA data type from an explicit value or the view's dtype."""
299 if data_type is not None:
300 if isinstance(data_type, TensorMapDataType):
301 return int(data_type)
302 dt = _normalize_tensor_map_data_type(data_type)
303 tma_dt = _NUMPY_DTYPE_TO_TMA.get(dt)
304 if tma_dt is None:
305 raise ValueError(
306 f"Unsupported dtype {dt} for TMA; "
307 f"supported dtypes: {list(_NUMPY_DTYPE_TO_TMA.keys())}.")
308 return tma_dt
310 dt = view.dtype
311 if dt is None:
312 raise ValueError(
313 "Cannot infer TMA data type from the tensor; "
314 "please specify data_type explicitly")
316 tma_dt = _NUMPY_DTYPE_TO_TMA.get(dt)
317 if tma_dt is None:
318 raise ValueError(
319 f"Unsupported dtype {dt} for TMA; "
320 f"supported dtypes: {list(_NUMPY_DTYPE_TO_TMA.keys())}. "
321 "You may also specify data_type explicitly.")
323 return tma_dt
326cdef inline bint _tma_dtype_to_dlpack(
327 int tma_dt,
328 uint8_t* out_code,
329 uint8_t* out_bits,
330 uint16_t* out_lanes,
331) noexcept:
332 if tma_dt == _TMA_DT_UINT8:
333 out_code[0] = <uint8_t>kDLUInt
334 out_bits[0] = <uint8_t>8
335 out_lanes[0] = <uint16_t>1
336 return True
337 if tma_dt == _TMA_DT_UINT16:
338 out_code[0] = <uint8_t>kDLUInt
339 out_bits[0] = <uint8_t>16
340 out_lanes[0] = <uint16_t>1
341 return True
342 if tma_dt == _TMA_DT_UINT32:
343 out_code[0] = <uint8_t>kDLUInt
344 out_bits[0] = <uint8_t>32
345 out_lanes[0] = <uint16_t>1
346 return True
347 if tma_dt == _TMA_DT_UINT64:
348 out_code[0] = <uint8_t>kDLUInt
349 out_bits[0] = <uint8_t>64
350 out_lanes[0] = <uint16_t>1
351 return True
352 if tma_dt == _TMA_DT_INT32:
353 out_code[0] = <uint8_t>kDLInt
354 out_bits[0] = <uint8_t>32
355 out_lanes[0] = <uint16_t>1
356 return True
357 if tma_dt == _TMA_DT_INT64:
358 out_code[0] = <uint8_t>kDLInt
359 out_bits[0] = <uint8_t>64
360 out_lanes[0] = <uint16_t>1
361 return True
362 if tma_dt == _TMA_DT_FLOAT16:
363 out_code[0] = <uint8_t>kDLFloat
364 out_bits[0] = <uint8_t>16
365 out_lanes[0] = <uint16_t>1
366 return True
367 if tma_dt == _TMA_DT_FLOAT32:
368 out_code[0] = <uint8_t>kDLFloat
369 out_bits[0] = <uint8_t>32
370 out_lanes[0] = <uint16_t>1
371 return True
372 if tma_dt == _TMA_DT_FLOAT64:
373 out_code[0] = <uint8_t>kDLFloat
374 out_bits[0] = <uint8_t>64
375 out_lanes[0] = <uint16_t>1
376 return True
377 if tma_dt == _TMA_DT_BFLOAT16:
378 out_code[0] = <uint8_t>kDLBfloat
379 out_bits[0] = <uint8_t>16
380 out_lanes[0] = <uint16_t>1
381 return True
382 return False
385cdef inline int _validate_tensor_map_view(view) except -1:
386 if not view.is_device_accessible: 1b
387 raise ValueError("The tensor must be device-accessible") 1b
389 if view.ptr % 16 != 0:
390 raise ValueError(
391 f"Global memory address must be 16-byte aligned, "
392 f"got address 0x{view.ptr:x}")
393 return 0
396def _get_validated_view(tensor):
397 """Obtain a device-accessible StridedMemoryView with a 16-byte-aligned pointer."""
398 if isinstance(tensor, StridedMemoryView):
399 view = tensor
400 else:
401 # stream_ptr=-1: no stream synchronization needed because descriptor
402 # creation only reads tensor metadata, it does not move data.
403 view = StridedMemoryView.from_any_interface(tensor, stream_ptr=-1)
404 _validate_tensor_map_view(view)
405 return view
408def _require_view_device(view, expected_device_id, operation):
409 """Ensure device-local tensors match the current CUDA device.
411 DLPack reports host/managed CUDA memory as ``kDLCUDAHost`` /
412 ``kDLCUDAManaged`` with ``device_id=0`` regardless of the current device,
413 so only true ``kDLCUDA`` tensors are rejected by device-id mismatch.
414 """
415 device_type, device_id = view.__dlpack_device__() 1defc
416 if device_type == _kDLCUDA and device_id != expected_device_id: 1defc
417 raise ValueError( 1c
418 f"{operation} expects tensor on device {expected_device_id}, got {device_id}") 1ac
419cdef inline intptr_t _get_current_context_ptr() except? 0:
420 cdef cydriver.CUcontext ctx
421 with nogil:
422 HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
423 if ctx == NULL:
424 raise RuntimeError("TensorMapDescriptor requires an active CUDA context")
425 return <intptr_t>ctx
428cdef inline int _get_current_device_id() except -1:
429 cdef cydriver.CUdevice dev
430 with nogil:
431 HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
432 return <int>dev
434def _compute_byte_strides(shape, strides, elem_size):
435 """Compute byte strides from element strides or C-contiguous fallback.
437 Returns a tuple of byte strides in row-major order.
438 """
439 if strides is not None:
440 return tuple(s * elem_size for s in strides)
442 # C-contiguous: compute byte strides from shape, innermost first
443 rank = len(shape)
444 byte_strides = []
445 stride = elem_size
446 for i in range(rank - 1, -1, -1):
447 byte_strides.append(stride)
448 stride *= shape[i]
449 byte_strides.reverse()
450 return tuple(byte_strides)
453def _validate_element_strides(element_strides, rank):
454 """Validate or default element_strides to all-ones."""
455 if element_strides is not None:
456 if len(element_strides) != rank:
457 raise ValueError(
458 f"element_strides must have {rank} elements, got {len(element_strides)}")
459 return element_strides
460 return (1,) * rank
463cdef class TensorMapDescriptor:
464 """Describes a TMA (Tensor Memory Accelerator) tensor map for Hopper+ GPUs.
466 A ``TensorMapDescriptor`` wraps the opaque 128-byte ``CUtensorMap`` struct
467 used by the hardware TMA unit for efficient bulk data movement between
468 global and shared memory.
470 Public tiled descriptors are created via
471 :meth:`cuda.core.StridedMemoryView.as_tensor_map`. Specialized
472 ``_from_*`` helpers remain private while this API surface settles, and
473 descriptors can be passed directly to :func:`~cuda.core.launch` as a
474 kernel argument.
475 """
477 def __init__(self):
478 raise RuntimeError( 1g
479 "TensorMapDescriptor cannot be instantiated directly. "
480 "Use StridedMemoryView.as_tensor_map() instead.")
482 cdef void* _get_data_ptr(self):
483 return <void*>&self._tensor_map
485 cdef int _check_context_compat(self) except -1:
486 cdef cydriver.CUcontext current_ctx
487 cdef cydriver.CUdevice current_dev
488 if self._context == 0 and self._device_id < 0:
489 return 0
490 with nogil:
491 HANDLE_RETURN(cydriver.cuCtxGetCurrent(¤t_ctx))
492 if current_ctx == NULL:
493 raise RuntimeError("TensorMapDescriptor requires an active CUDA context")
494 if self._context != 0 and <intptr_t>current_ctx != self._context:
495 raise RuntimeError(
496 "TensorMapDescriptor was created in a different CUDA context")
497 with nogil:
498 HANDLE_RETURN(cydriver.cuCtxGetDevice(¤t_dev))
499 cdef int current_dev_id = <int>current_dev
500 if self._device_id >= 0 and current_dev_id != self._device_id:
501 raise RuntimeError(
502 f"TensorMapDescriptor belongs to device {self._device_id}, "
503 f"but current device is {current_dev_id}")
504 return 0
506 @property
507 def device(self) -> Device | None:
508 """Return the :obj:`~cuda.core.Device` associated with this descriptor."""
509 if self._device_id >= 0:
510 from cuda.core._device import Device
511 return Device(self._device_id)
512 return None
514 @classmethod
515 def _from_tiled(cls, view, box_dim=None, *,
516 options=None,
517 element_strides=None,
518 data_type=None,
519 interleave=TensorMapInterleave.NONE,
520 swizzle=TensorMapSwizzle.NONE,
521 l2_promotion=TensorMapL2Promotion.NONE,
522 oob_fill=TensorMapOOBFill.NONE):
523 """Create a tiled TMA descriptor from a validated view.
525 Parameters
526 ----------
527 view : StridedMemoryView
528 A device-accessible view with a 16-byte-aligned pointer.
529 box_dim : tuple of int, optional
530 The size of each tile dimension (in elements). Must have the
531 same rank as the tensor and each value must be in [1, 256].
532 Specified in the same (row-major) order as the tensor shape.
533 Required unless ``options`` is provided.
534 options : TensorMapDescriptorOptions or mapping, optional
535 Bundled tiled-descriptor options. When provided, do not also pass
536 ``box_dim`` or the individual option kwargs.
537 element_strides : tuple of int, optional
538 Per-dimension element traversal strides. Default is all 1s.
539 Specified in the same (row-major) order as the tensor shape.
540 data_type : dtype-like or TensorMapDataType, optional
541 Explicit dtype override. If ``None``, inferred from the tensor's
542 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
543 accepted for compatibility.
544 interleave : TensorMapInterleave
545 Interleave layout. Default ``NONE``.
546 swizzle : TensorMapSwizzle
547 Swizzle mode. Default ``NONE``.
548 l2_promotion : TensorMapL2Promotion
549 L2 promotion mode. Default ``NONE``.
550 oob_fill : TensorMapOOBFill
551 Out-of-bounds fill mode. Default ``NONE``.
553 Returns
554 -------
555 TensorMapDescriptor
557 Raises
558 ------
559 ValueError
560 If the tensor rank is outside [1, 5], the pointer is not
561 16-byte aligned, or dimension/stride constraints are violated.
562 """
563 cdef TensorMapDescriptor desc = cls.__new__(cls) 1b
565 opts = _coerce_tensor_map_descriptor_options( 1ab
566 box_dim,
567 options,
568 element_strides=element_strides,
569 data_type=data_type,
570 interleave=interleave,
571 swizzle=swizzle,
572 l2_promotion=l2_promotion,
573 oob_fill=oob_fill, 1b
574 )
575 box_dim = opts.box_dim 1b
576 element_strides = opts.element_strides 1b
577 data_type = opts.data_type 1b
578 interleave = opts.interleave 1b
579 swizzle = opts.swizzle 1b
580 l2_promotion = opts.l2_promotion 1b
581 oob_fill = opts.oob_fill 1b
583 _validate_tensor_map_view(view) 1ab
584 # Keep both the original tensor object and the validated view alive.
585 # For DLPack exporters, the view may hold the owning capsule whose
586 # deleter can free the backing allocation when released.
587 desc._source_ref = view.exporting_obj
588 desc._view_ref = view
589 desc._context = _get_current_context_ptr()
590 desc._device_id = _get_current_device_id()
591 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_tiled")
593 tma_dt = _resolve_data_type(view, data_type)
594 cdef int c_data_type_int = tma_dt
595 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int
597 cdef intptr_t global_address = view.ptr
598 shape = view.shape
600 cdef int rank = len(shape)
601 if rank < 1 or rank > 5:
602 raise ValueError(
603 f"Tensor rank must be between 1 and 5, got {rank}")
605 if len(box_dim) != rank:
606 raise ValueError(
607 f"box_dim must have {rank} elements (same as tensor rank), "
608 f"got {len(box_dim)}")
610 for i, bd in enumerate(box_dim):
611 if bd < 1 or bd > 256:
612 raise ValueError(
613 f"box_dim[{i}] must be in [1, 256], got {bd}")
615 cdef bint elem_strides_provided = element_strides is not None
616 element_strides = _validate_element_strides(element_strides, rank)
618 # Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible.
619 # This avoids maintaining a second, independent validation/encoding implementation.
620 cdef uint8_t dl_code
621 cdef uint8_t dl_bits
622 cdef uint16_t dl_lanes
623 cdef int64_t c_shape[5]
624 cdef int64_t c_strides[5]
625 cdef int c_box_sizes[5]
626 cdef int c_elem_strides[5]
627 cdef const int64_t* c_strides_ptr
628 cdef const int* c_elem_strides_ptr
629 cdef char errbuf[512]
630 cdef int i_cccl
631 cdef int device_type
632 cdef int c_device_id
633 cdef int dl_device_type
634 cdef int dl_device_id
635 cdef int c_cccl_interleave_int
636 cdef int c_cccl_swizzle_int
637 cdef int c_cccl_l2_promotion_int
638 cdef int c_cccl_oob_fill_int
639 cdef int rc
640 if _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes):
641 c_strides_ptr = NULL
642 c_elem_strides_ptr = NULL
643 errbuf[0] = 0
645 for i_cccl in range(rank):
646 c_shape[i_cccl] = <int64_t>shape[i_cccl]
647 c_box_sizes[i_cccl] = <int>box_dim[i_cccl]
648 if elem_strides_provided:
649 c_elem_strides[i_cccl] = <int>element_strides[i_cccl]
651 if view.strides is not None:
652 for i_cccl in range(rank):
653 c_strides[i_cccl] = <int64_t>view.strides[i_cccl]
654 c_strides_ptr = &c_strides[0]
656 if elem_strides_provided:
657 c_elem_strides_ptr = &c_elem_strides[0]
659 dl_device_type, dl_device_id = view.__dlpack_device__()
660 device_type = dl_device_type
661 c_device_id = dl_device_id
662 c_cccl_interleave_int = int(interleave)
663 c_cccl_swizzle_int = int(swizzle)
664 c_cccl_l2_promotion_int = int(l2_promotion)
665 c_cccl_oob_fill_int = int(oob_fill)
667 with nogil:
668 rc = cuda_core_cccl_make_tma_descriptor_tiled(
669 <void*>&desc._tensor_map,
670 <void*>global_address,
671 device_type,
672 c_device_id,
673 rank,
674 &c_shape[0],
675 c_strides_ptr,
676 dl_code,
677 dl_bits,
678 dl_lanes,
679 &c_box_sizes[0],
680 c_elem_strides_ptr,
681 c_cccl_interleave_int,
682 c_cccl_swizzle_int,
683 c_cccl_l2_promotion_int,
684 c_cccl_oob_fill_int,
685 &errbuf[0],
686 <size_t>sizeof(errbuf),
687 )
689 if rc == 0:
690 desc._repr_info = {
691 "method": "tiled",
692 "rank": rank,
693 "data_type": TensorMapDataType(tma_dt),
694 "swizzle": swizzle,
695 }
696 return desc
698 msg = errbuf[:].split(b"\0", 1)[0].decode("utf-8", errors="replace")
699 # If CCCL isn't available at build time, fall back to the direct
700 # driver API path to preserve functionality on older toolchains.
701 if "not available at build time" not in msg:
702 raise ValueError(f"Failed to build TMA descriptor via CCCL: {msg}")
704 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
705 byte_strides = _compute_byte_strides(shape, view.strides, elem_size)
707 # Reverse dimensions for column-major cuTensorMap convention
708 # Python/DLPack: row-major (dim 0 = outermost)
709 # cuTensorMap: column-major (dim 0 = innermost)
710 cdef uint64_t[5] c_global_dim
711 cdef uint64_t[4] c_global_strides # rank - 1 elements
712 cdef uint32_t[5] c_box_dim
713 cdef uint32_t[5] c_element_strides
714 cdef int i_c
716 for i_c in range(rank):
717 # Reverse: Python dim i -> cuTensorMap dim (rank - 1 - i)
718 c_global_dim[i_c] = <uint64_t>shape[rank - 1 - i_c]
719 c_box_dim[i_c] = <uint32_t>box_dim[rank - 1 - i_c]
720 c_element_strides[i_c] = <uint32_t>element_strides[rank - 1 - i_c]
722 # globalStrides: rank-1 elements (byte strides for dims 1..N-1 in col-major order)
723 # The innermost stride (dim 0) is implicit = element size
724 for i_c in range(rank - 1):
725 c_global_strides[i_c] = <uint64_t>byte_strides[rank - 2 - i_c]
727 cdef uint32_t c_rank = <uint32_t>rank
728 cdef int c_interleave_int = int(interleave)
729 cdef int c_swizzle_int = int(swizzle)
730 cdef int c_l2_promotion_int = int(l2_promotion)
731 cdef int c_oob_fill_int = int(oob_fill)
732 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave>c_interleave_int
733 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle>c_swizzle_int
734 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion>c_l2_promotion_int
735 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill>c_oob_fill_int
737 with nogil:
738 HANDLE_RETURN(cydriver.cuTensorMapEncodeTiled(
739 &desc._tensor_map,
740 c_data_type,
741 c_rank,
742 <void*>global_address,
743 c_global_dim,
744 c_global_strides,
745 c_box_dim,
746 c_element_strides,
747 c_interleave,
748 c_swizzle,
749 c_l2_promotion,
750 c_oob_fill,
751 ))
753 desc._repr_info = {
754 "method": "tiled",
755 "rank": rank,
756 "data_type": TensorMapDataType(tma_dt),
757 "swizzle": swizzle,
758 }
760 return desc
762 @classmethod
763 def _from_im2col(cls, view, pixel_box_lower_corner, pixel_box_upper_corner,
764 channels_per_pixel, pixels_per_column, *,
765 element_strides=None,
766 data_type=None,
767 interleave=TensorMapInterleave.NONE,
768 swizzle=TensorMapSwizzle.NONE,
769 l2_promotion=TensorMapL2Promotion.NONE,
770 oob_fill=TensorMapOOBFill.NONE):
771 """Create an im2col TMA descriptor from a validated view.
773 Im2col layout is used for convolution-style data access patterns.
775 Parameters
776 ----------
777 view : StridedMemoryView
778 A device-accessible view with a 16-byte-aligned pointer.
779 pixel_box_lower_corner : tuple of int
780 Lower corner of the pixel bounding box for each spatial
781 dimension (rank - 2 elements). Specified in row-major order
782 matching the tensor's spatial dimensions.
783 pixel_box_upper_corner : tuple of int
784 Upper corner of the pixel bounding box for each spatial
785 dimension (rank - 2 elements). Specified in row-major order
786 matching the tensor's spatial dimensions.
787 channels_per_pixel : int
788 Number of channels per pixel.
789 pixels_per_column : int
790 Number of pixels per column.
791 element_strides : tuple of int, optional
792 Per-dimension element traversal strides. Default is all 1s.
793 data_type : dtype-like or TensorMapDataType, optional
794 Explicit dtype override. If ``None``, inferred from the tensor's
795 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
796 accepted for compatibility.
797 interleave : TensorMapInterleave
798 Interleave layout. Default ``NONE``.
799 swizzle : TensorMapSwizzle
800 Swizzle mode. Default ``NONE``.
801 l2_promotion : TensorMapL2Promotion
802 L2 promotion mode. Default ``NONE``.
803 oob_fill : TensorMapOOBFill
804 Out-of-bounds fill mode. Default ``NONE``.
806 Returns
807 -------
808 TensorMapDescriptor
810 Raises
811 ------
812 ValueError
813 If the tensor rank is outside [3, 5], the pointer is not
814 16-byte aligned, or other constraints are violated.
815 """
816 cdef TensorMapDescriptor desc = cls.__new__(cls)
818 _validate_tensor_map_view(view)
819 desc._source_ref = view.exporting_obj
820 desc._view_ref = view
821 desc._context = _get_current_context_ptr()
822 desc._device_id = _get_current_device_id()
823 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_im2col")
825 tma_dt = _resolve_data_type(view, data_type)
826 cdef int c_data_type_int = tma_dt
827 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int
829 cdef intptr_t global_address = view.ptr
830 shape = view.shape
832 cdef int rank = len(shape)
833 if rank < 3 or rank > 5:
834 raise ValueError(
835 f"Im2col tensor rank must be between 3 and 5, got {rank}")
837 cdef int n_spatial = rank - 2
838 if len(pixel_box_lower_corner) != n_spatial:
839 raise ValueError(
840 f"pixel_box_lower_corner must have {n_spatial} elements "
841 f"(rank - 2), got {len(pixel_box_lower_corner)}")
842 if len(pixel_box_upper_corner) != n_spatial:
843 raise ValueError(
844 f"pixel_box_upper_corner must have {n_spatial} elements "
845 f"(rank - 2), got {len(pixel_box_upper_corner)}")
847 element_strides = _validate_element_strides(element_strides, rank)
849 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
850 byte_strides = _compute_byte_strides(shape, view.strides, elem_size)
852 # Reverse all dimension arrays for column-major convention
853 cdef uint64_t[5] c_global_dim
854 cdef uint64_t[4] c_global_strides
855 cdef uint32_t[5] c_element_strides
856 cdef int[3] c_pixel_box_lower # max 3 spatial dims (rank 5 - 2)
857 cdef int[3] c_pixel_box_upper
858 cdef int i_c
860 for i_c in range(3):
861 c_pixel_box_lower[i_c] = 0
862 c_pixel_box_upper[i_c] = 0
864 for i_c in range(rank):
865 c_global_dim[i_c] = <uint64_t>shape[rank - 1 - i_c]
866 c_element_strides[i_c] = <uint32_t>element_strides[rank - 1 - i_c]
868 for i_c in range(rank - 1):
869 c_global_strides[i_c] = <uint64_t>byte_strides[rank - 2 - i_c]
871 # Reverse spatial dimensions for lower/upper corners
872 for i_c in range(n_spatial):
873 c_pixel_box_lower[i_c] = <int>pixel_box_lower_corner[n_spatial - 1 - i_c]
874 c_pixel_box_upper[i_c] = <int>pixel_box_upper_corner[n_spatial - 1 - i_c]
876 cdef uint32_t c_rank = <uint32_t>rank
877 cdef uint32_t c_channels = <uint32_t>channels_per_pixel
878 cdef uint32_t c_pixels = <uint32_t>pixels_per_column
879 cdef int c_interleave_int = int(interleave)
880 cdef int c_swizzle_int = int(swizzle)
881 cdef int c_l2_promotion_int = int(l2_promotion)
882 cdef int c_oob_fill_int = int(oob_fill)
883 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave>c_interleave_int
884 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle>c_swizzle_int
885 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion>c_l2_promotion_int
886 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill>c_oob_fill_int
888 with nogil:
889 HANDLE_RETURN(cydriver.cuTensorMapEncodeIm2col(
890 &desc._tensor_map,
891 c_data_type,
892 c_rank,
893 <void*>global_address,
894 c_global_dim,
895 c_global_strides,
896 c_pixel_box_lower,
897 c_pixel_box_upper,
898 c_channels,
899 c_pixels,
900 c_element_strides,
901 c_interleave,
902 c_swizzle,
903 c_l2_promotion,
904 c_oob_fill,
905 ))
907 desc._repr_info = {
908 "method": "im2col",
909 "rank": rank,
910 "data_type": TensorMapDataType(tma_dt),
911 "swizzle": swizzle,
912 }
914 return desc
916 @classmethod
917 def _from_im2col_wide(cls, view, pixel_box_lower_corner_width, pixel_box_upper_corner_width,
918 channels_per_pixel, pixels_per_column, *,
919 element_strides=None,
920 data_type=None,
921 interleave=TensorMapInterleave.NONE,
922 mode=TensorMapIm2ColWideMode.W,
923 swizzle=TensorMapSwizzle.SWIZZLE_128B,
924 l2_promotion=TensorMapL2Promotion.NONE,
925 oob_fill=TensorMapOOBFill.NONE):
926 """Create an im2col-wide TMA descriptor from a validated view.
928 Im2col-wide layout loads elements exclusively along the W (width)
929 dimension. This variant is supported on compute capability 10.0+
930 (Blackwell and later).
932 Parameters
933 ----------
934 view : StridedMemoryView
935 A device-accessible view with a 16-byte-aligned pointer.
936 pixel_box_lower_corner_width : int
937 Lower corner of the pixel bounding box along the W dimension.
938 pixel_box_upper_corner_width : int
939 Upper corner of the pixel bounding box along the W dimension.
940 channels_per_pixel : int
941 Number of channels per pixel.
942 pixels_per_column : int
943 Number of pixels per column.
944 element_strides : tuple of int, optional
945 Per-dimension element traversal strides. Default is all 1s.
946 data_type : dtype-like or TensorMapDataType, optional
947 Explicit dtype override. If ``None``, inferred from the tensor's
948 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
949 accepted for compatibility.
950 interleave : TensorMapInterleave
951 Interleave layout. Default ``NONE``.
952 mode : TensorMapIm2ColWideMode
953 Im2col wide mode. Default ``W``.
954 swizzle : TensorMapSwizzle
955 Swizzle mode. Default ``SWIZZLE_128B``.
956 l2_promotion : TensorMapL2Promotion
957 L2 promotion mode. Default ``NONE``.
958 oob_fill : TensorMapOOBFill
959 Out-of-bounds fill mode. Default ``NONE``.
961 Returns
962 -------
963 TensorMapDescriptor
965 Raises
966 ------
967 ValueError
968 If the tensor rank is outside [3, 5], the pointer is not
969 16-byte aligned, or other constraints are violated.
970 """
971 IF CUDA_CORE_BUILD_MAJOR < 13:
972 raise RuntimeError(
973 "TensorMapDescriptor._from_im2col_wide requires a CUDA 13+ build")
974 ELSE:
975 cdef TensorMapDescriptor desc = cls.__new__(cls)
977 _validate_tensor_map_view(view)
978 desc._source_ref = view.exporting_obj
979 desc._view_ref = view
980 desc._context = _get_current_context_ptr()
981 desc._device_id = _get_current_device_id()
982 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_im2col_wide")
984 tma_dt = _resolve_data_type(view, data_type)
985 cdef int c_data_type_int = tma_dt
986 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int
988 cdef intptr_t global_address = view.ptr
989 shape = view.shape
991 cdef int rank = len(shape)
992 if rank < 3 or rank > 5:
993 raise ValueError(
994 f"Im2col-wide tensor rank must be between 3 and 5, got {rank}")
996 element_strides = _validate_element_strides(element_strides, rank)
998 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
999 byte_strides = _compute_byte_strides(shape, view.strides, elem_size)
1001 # Reverse all dimension arrays for column-major convention
1002 cdef uint64_t[5] c_global_dim
1003 cdef uint64_t[4] c_global_strides
1004 cdef uint32_t[5] c_element_strides
1005 cdef int i_c
1007 for i_c in range(rank):
1008 c_global_dim[i_c] = <uint64_t>shape[rank - 1 - i_c]
1009 c_element_strides[i_c] = <uint32_t>element_strides[rank - 1 - i_c]
1011 for i_c in range(rank - 1):
1012 c_global_strides[i_c] = <uint64_t>byte_strides[rank - 2 - i_c]
1014 cdef uint32_t c_rank = <uint32_t>rank
1015 cdef int c_lower_w = <int>pixel_box_lower_corner_width
1016 cdef int c_upper_w = <int>pixel_box_upper_corner_width
1017 cdef uint32_t c_channels = <uint32_t>channels_per_pixel
1018 cdef uint32_t c_pixels = <uint32_t>pixels_per_column
1019 cdef int c_interleave_int = int(interleave)
1020 cdef int c_mode_int = int(mode)
1021 cdef int c_swizzle_int = int(swizzle)
1022 cdef int c_l2_promotion_int = int(l2_promotion)
1023 cdef int c_oob_fill_int = int(oob_fill)
1024 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave>c_interleave_int
1025 cdef cydriver.CUtensorMapIm2ColWideMode c_mode = <cydriver.CUtensorMapIm2ColWideMode>c_mode_int
1026 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle>c_swizzle_int
1027 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion>c_l2_promotion_int
1028 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill>c_oob_fill_int
1030 with nogil:
1031 HANDLE_RETURN(cydriver.cuTensorMapEncodeIm2colWide(
1032 &desc._tensor_map,
1033 c_data_type,
1034 c_rank,
1035 <void*>global_address,
1036 c_global_dim,
1037 c_global_strides,
1038 c_lower_w,
1039 c_upper_w,
1040 c_channels,
1041 c_pixels,
1042 c_element_strides,
1043 c_interleave,
1044 c_mode,
1045 c_swizzle,
1046 c_l2_promotion,
1047 c_oob_fill,
1048 ))
1050 desc._repr_info = {
1051 "method": "im2col_wide",
1052 "rank": rank,
1053 "data_type": TensorMapDataType(tma_dt),
1054 "swizzle": swizzle,
1055 }
1057 return desc
1059 def replace_address(self, tensor: object) -> None:
1060 """Replace the global memory address in this tensor map descriptor.
1062 This is useful when the tensor data has been reallocated but the
1063 shape, strides, and other parameters remain the same.
1065 Parameters
1066 ----------
1067 tensor : object
1068 Any object supporting DLPack or ``__cuda_array_interface__``,
1069 or a :obj:`~cuda.core.StridedMemoryView`. Must refer to
1070 device-accessible memory with a 16-byte-aligned pointer.
1071 """
1072 self._check_context_compat()
1073 view = _get_validated_view(tensor)
1074 _require_view_device(view, self._device_id, "replace_address")
1076 cdef intptr_t global_address = view.ptr
1078 with nogil:
1079 HANDLE_RETURN(cydriver.cuTensorMapReplaceAddress(
1080 &self._tensor_map,
1081 <void*>global_address,
1082 ))
1084 # Update the source reference only after the driver call succeeds,
1085 # so we don't drop the old tensor (risking a dangling pointer in the
1086 # CUtensorMap struct) if the call fails.
1087 self._source_ref = view.exporting_obj
1088 self._view_ref = view
1090 def __repr__(self) -> str:
1091 info = self._repr_info
1092 if info is None:
1093 return "TensorMapDescriptor()"
1094 parts = []
1095 if "method" in info:
1096 parts.append(info["method"])
1097 if "rank" in info:
1098 parts.append(f"rank={info['rank']}")
1099 if "data_type" in info:
1100 parts.append(f"dtype={info['data_type'].name}")
1101 if "swizzle" in info:
1102 parts.append(f"swizzle={info['swizzle'].name}")
1103 return f"TensorMapDescriptor({', '.join(parts)})"