Coverage for cuda/core/_tensor_map.pyx: 39.96%
528 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 02:41 +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
50__all__ = ['TensorMapDescriptor', 'TensorMapDescriptorOptions']
53class TensorMapDataType(enum.IntEnum):
54 """Data types for tensor map descriptors.
56 These correspond to the ``CUtensorMapDataType`` driver enum values.
57 """
58 UINT8 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8
59 UINT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16
60 UINT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32
61 INT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32
62 UINT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64
63 INT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64
64 FLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16
65 FLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32
66 FLOAT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64
67 BFLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16
68 FLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ
69 TFLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32
70 TFLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ
73class TensorMapInterleave(enum.IntEnum):
74 """Interleave layout for tensor map descriptors.
76 These correspond to the ``CUtensorMapInterleave`` driver enum values.
77 """
78 NONE = cydriver.CU_TENSOR_MAP_INTERLEAVE_NONE
79 INTERLEAVE_16B = cydriver.CU_TENSOR_MAP_INTERLEAVE_16B
80 INTERLEAVE_32B = cydriver.CU_TENSOR_MAP_INTERLEAVE_32B
83class TensorMapSwizzle(enum.IntEnum):
84 """Swizzle mode for tensor map descriptors.
86 These correspond to the ``CUtensorMapSwizzle`` driver enum values.
87 """
88 NONE = cydriver.CU_TENSOR_MAP_SWIZZLE_NONE
89 SWIZZLE_32B = cydriver.CU_TENSOR_MAP_SWIZZLE_32B
90 SWIZZLE_64B = cydriver.CU_TENSOR_MAP_SWIZZLE_64B
91 SWIZZLE_128B = cydriver.CU_TENSOR_MAP_SWIZZLE_128B
94class TensorMapL2Promotion(enum.IntEnum):
95 """L2 promotion mode for tensor map descriptors.
97 These correspond to the ``CUtensorMapL2promotion`` driver enum values.
98 """
99 NONE = cydriver.CU_TENSOR_MAP_L2_PROMOTION_NONE
100 L2_64B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_64B
101 L2_128B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_128B
102 L2_256B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_256B
105class TensorMapOOBFill(enum.IntEnum):
106 """Out-of-bounds fill mode for tensor map descriptors.
108 These correspond to the ``CUtensorMapFloatOOBfill`` driver enum values.
109 """
110 NONE = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
111 NAN_REQUEST_ZERO_FMA = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA
114IF CUDA_CORE_BUILD_MAJOR >= 13:
115 class TensorMapIm2ColWideMode(enum.IntEnum):
116 """Im2col wide mode for tensor map descriptors.
118 These correspond to the ``CUtensorMapIm2ColWideMode`` driver enum values.
119 Supported on compute capability 10.0+.
120 """
121 W = cydriver.CU_TENSOR_MAP_IM2COL_WIDE_MODE_W
122 W128 = cydriver.CU_TENSOR_MAP_IM2COL_WIDE_MODE_W128
123ELSE:
124 class TensorMapIm2ColWideMode(enum.IntEnum):
125 """Im2col wide mode for tensor map descriptors.
127 This enum is always defined for API stability, but the
128 :meth:`TensorMapDescriptor._from_im2col_wide` factory requires a CUDA 13+
129 build and will raise otherwise.
130 """
131 W = 0
132 W128 = 1
135_TMA_DT_UINT8: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8)
136_TMA_DT_UINT16: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16)
137_TMA_DT_UINT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32)
138_TMA_DT_INT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32)
139_TMA_DT_UINT64: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64)
140_TMA_DT_INT64: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64)
141_TMA_DT_FLOAT16: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16)
142_TMA_DT_FLOAT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32)
143_TMA_DT_FLOAT64: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64)
144_TMA_DT_BFLOAT16: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16)
145_TMA_DT_FLOAT32_FTZ: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ)
146_TMA_DT_TFLOAT32: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32)
147_TMA_DT_TFLOAT32_FTZ: int = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ)
150def _normalize_tensor_map_data_type(data_type):
151 if data_type is None or isinstance(data_type, TensorMapDataType): 1bdce
152 return data_type 1bdc
153 try: 1e
154 return numpy.dtype(data_type) 1ae
155 except TypeError as e:
156 raise TypeError(
157 "data_type must be a TensorMapDataType or a numpy/ml_dtypes dtype, "
158 f"got {type(data_type)}") from e
161def _normalize_tensor_map_sequence(name, values):
162 try: 1bghdc
163 values = tuple(values) 1bghdc
164 except TypeError as e: 1ah
165 raise TypeError(f"{name} must be a tuple of ints, got {type(values)}") from e 1h
166 for i, value in enumerate(values): 1bgdc
167 if not isinstance(value, int): 1bgdc
168 raise TypeError(f"{name}[{i}] must be an int, got {type(value)}") 1g
169 return values 1bdc
172def _require_tensor_map_enum(name, value, enum_type):
173 if not isinstance(value, enum_type): 1bdc
174 raise TypeError(f"{name} must be a {enum_type.__name__}, got {type(value)}") 1dc
175 return value 1bc
178@dataclass
179class TensorMapDescriptorOptions:
180 """Options for :meth:`cuda.core.StridedMemoryView.as_tensor_map`.
182 Attributes
183 ----------
184 box_dim : tuple[int, ...]
185 Tile size for each tensor dimension, expressed in elements.
186 element_strides : tuple[int, ...], optional
187 Per-dimension element traversal strides.
188 data_type : object, optional
189 Explicit dtype override. Prefer NumPy or ``ml_dtypes`` dtype objects;
190 :class:`TensorMapDataType` remains accepted for compatibility.
191 interleave : TensorMapInterleave, optional
192 Interleave layout. Default ``NONE``.
193 swizzle : TensorMapSwizzle, optional
194 Swizzle mode. Default ``NONE``.
195 l2_promotion : TensorMapL2Promotion, optional
196 L2 promotion mode. Default ``NONE``.
197 oob_fill : TensorMapOOBFill, optional
198 Out-of-bounds fill mode. Default ``NONE``.
199 """
201 box_dim: tuple[int, ...]
202 element_strides: tuple[int, ...] | None = None
203 data_type: object = None
204 interleave: TensorMapInterleave = TensorMapInterleave.NONE
205 swizzle: TensorMapSwizzle = TensorMapSwizzle.NONE
206 l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE
207 oob_fill: TensorMapOOBFill = TensorMapOOBFill.NONE
209 def __post_init__(self) -> None:
210 self.box_dim = _normalize_tensor_map_sequence("box_dim", self.box_dim) 1bghdc
211 if self.element_strides is not None: 1bdc
212 self.element_strides = _normalize_tensor_map_sequence("element_strides", self.element_strides) 1b
213 self.data_type = _normalize_tensor_map_data_type(self.data_type) 1abdc
214 self.interleave = _require_tensor_map_enum("interleave", self.interleave, TensorMapInterleave) 1bdc
215 self.swizzle = _require_tensor_map_enum("swizzle", self.swizzle, TensorMapSwizzle) 1bc
216 self.l2_promotion = _require_tensor_map_enum("l2_promotion", self.l2_promotion, TensorMapL2Promotion) 1b
217 self.oob_fill = _require_tensor_map_enum("oob_fill", self.oob_fill, TensorMapOOBFill) 1b
220def _coerce_tensor_map_descriptor_options(
221 box_dim,
222 options,
223 *,
224 element_strides,
225 data_type,
226 interleave,
227 swizzle,
228 l2_promotion,
229 oob_fill,
230):
231 if options is not None: 1bk
232 if (
233 box_dim is not None
234 or element_strides is not None
235 or data_type is not None
236 or interleave != TensorMapInterleave.NONE
237 or swizzle != TensorMapSwizzle.NONE
238 or l2_promotion != TensorMapL2Promotion.NONE
239 or oob_fill != TensorMapOOBFill.NONE
240 ):
241 raise TypeError(
242 "Specify either options or the individual tensor map arguments, not both")
243 return check_or_create_options(
244 TensorMapDescriptorOptions,
245 options,
246 "Tensor map descriptor options",
247 )
249 if box_dim is None: 1bk
250 raise TypeError("box_dim is required unless options is provided") 1k
252 return TensorMapDescriptorOptions( 1b
253 box_dim=box_dim,
254 element_strides=element_strides,
255 data_type=data_type,
256 interleave=interleave,
257 swizzle=swizzle,
258 l2_promotion=l2_promotion,
259 oob_fill=oob_fill, 1b
260 )
263# Mapping from numpy dtype to TMA data type
264_NUMPY_DTYPE_TO_TMA = {
265 numpy.dtype(numpy.uint8): _TMA_DT_UINT8,
266 numpy.dtype(numpy.uint16): _TMA_DT_UINT16,
267 numpy.dtype(numpy.uint32): _TMA_DT_UINT32,
268 numpy.dtype(numpy.int32): _TMA_DT_INT32,
269 numpy.dtype(numpy.uint64): _TMA_DT_UINT64,
270 numpy.dtype(numpy.int64): _TMA_DT_INT64,
271 numpy.dtype(numpy.float16): _TMA_DT_FLOAT16,
272 numpy.dtype(numpy.float32): _TMA_DT_FLOAT32,
273 numpy.dtype(numpy.float64): _TMA_DT_FLOAT64,
274}
276if ml_bfloat16 is not None:
277 _NUMPY_DTYPE_TO_TMA[numpy.dtype(ml_bfloat16)] = _TMA_DT_BFLOAT16
280# Mapping from TMA data type to element size in bytes
281_TMA_DATA_TYPE_SIZE = {
282 _TMA_DT_UINT8: 1,
283 _TMA_DT_UINT16: 2,
284 _TMA_DT_UINT32: 4,
285 _TMA_DT_INT32: 4,
286 _TMA_DT_UINT64: 8,
287 _TMA_DT_INT64: 8,
288 _TMA_DT_FLOAT16: 2,
289 _TMA_DT_FLOAT32: 4,
290 _TMA_DT_FLOAT64: 8,
291 _TMA_DT_BFLOAT16: 2,
292 _TMA_DT_FLOAT32_FTZ: 4,
293 _TMA_DT_TFLOAT32: 4,
294 _TMA_DT_TFLOAT32_FTZ: 4,
295}
298def _resolve_data_type(view, data_type):
299 """Resolve the TMA data type from an explicit value or the view's dtype."""
301 if data_type is not None: 1ief
302 if isinstance(data_type, TensorMapDataType): 1e
303 return int(data_type)
304 dt = _normalize_tensor_map_data_type(data_type) 1e
305 tma_dt = _NUMPY_DTYPE_TO_TMA.get(dt) 1e
306 if tma_dt is None: 1e
307 raise ValueError( 1e
308 f"Unsupported dtype {dt} for TMA; " 1e
309 f"supported dtypes: {list(_NUMPY_DTYPE_TO_TMA.keys())}.") 1e
310 return tma_dt
312 dt = view.dtype 1if
313 if dt is None: 1if
314 raise ValueError( 1i
315 "Cannot infer TMA data type from the tensor; "
316 "please specify data_type explicitly")
318 tma_dt = _NUMPY_DTYPE_TO_TMA.get(dt) 1f
319 if tma_dt is None: 1f
320 raise ValueError( 1f
321 f"Unsupported dtype {dt} for TMA; " 1f
322 f"supported dtypes: {list(_NUMPY_DTYPE_TO_TMA.keys())}. " 1f
323 "You may also specify data_type explicitly.")
325 return tma_dt
328cdef inline bint _tma_dtype_to_dlpack(
329 int tma_dt,
330 uint8_t* out_code,
331 uint8_t* out_bits,
332 uint16_t* out_lanes,
333) noexcept:
334 if tma_dt == _TMA_DT_UINT8:
335 out_code[0] = <uint8_t>kDLUInt
336 out_bits[0] = <uint8_t>8
337 out_lanes[0] = <uint16_t>1
338 return True
339 if tma_dt == _TMA_DT_UINT16:
340 out_code[0] = <uint8_t>kDLUInt
341 out_bits[0] = <uint8_t>16
342 out_lanes[0] = <uint16_t>1
343 return True
344 if tma_dt == _TMA_DT_UINT32:
345 out_code[0] = <uint8_t>kDLUInt
346 out_bits[0] = <uint8_t>32
347 out_lanes[0] = <uint16_t>1
348 return True
349 if tma_dt == _TMA_DT_UINT64:
350 out_code[0] = <uint8_t>kDLUInt
351 out_bits[0] = <uint8_t>64
352 out_lanes[0] = <uint16_t>1
353 return True
354 if tma_dt == _TMA_DT_INT32:
355 out_code[0] = <uint8_t>kDLInt
356 out_bits[0] = <uint8_t>32
357 out_lanes[0] = <uint16_t>1
358 return True
359 if tma_dt == _TMA_DT_INT64:
360 out_code[0] = <uint8_t>kDLInt
361 out_bits[0] = <uint8_t>64
362 out_lanes[0] = <uint16_t>1
363 return True
364 if tma_dt == _TMA_DT_FLOAT16:
365 out_code[0] = <uint8_t>kDLFloat
366 out_bits[0] = <uint8_t>16
367 out_lanes[0] = <uint16_t>1
368 return True
369 if tma_dt == _TMA_DT_FLOAT32:
370 out_code[0] = <uint8_t>kDLFloat
371 out_bits[0] = <uint8_t>32
372 out_lanes[0] = <uint16_t>1
373 return True
374 if tma_dt == _TMA_DT_FLOAT64:
375 out_code[0] = <uint8_t>kDLFloat
376 out_bits[0] = <uint8_t>64
377 out_lanes[0] = <uint16_t>1
378 return True
379 if tma_dt == _TMA_DT_BFLOAT16:
380 out_code[0] = <uint8_t>kDLBfloat
381 out_bits[0] = <uint8_t>16
382 out_lanes[0] = <uint16_t>1
383 return True
384 return False
387cdef inline int _validate_tensor_map_view(view) except -1:
388 if not view.is_device_accessible: 1b
389 raise ValueError("The tensor must be device-accessible") 1b
391 if view.ptr % 16 != 0:
392 raise ValueError(
393 f"Global memory address must be 16-byte aligned, "
394 f"got address 0x{view.ptr:x}")
395 return 0
398def _get_validated_view(tensor):
399 """Obtain a device-accessible StridedMemoryView with a 16-byte-aligned pointer."""
400 if isinstance(tensor, StridedMemoryView):
401 view = tensor
402 else:
403 # stream_ptr=-1: no stream synchronization needed because descriptor
404 # creation only reads tensor metadata, it does not move data.
405 view = StridedMemoryView.from_any_interface(tensor, stream_ptr=-1)
406 _validate_tensor_map_view(view)
407 return view
410def _require_view_device(view, expected_device_id, operation):
411 """Ensure device-local tensors match the current CUDA device.
413 DLPack reports host/managed CUDA memory as ``kDLCUDAHost`` /
414 ``kDLCUDAManaged`` with ``device_id=0`` regardless of the current device,
415 so only true ``kDLCUDA`` tensors are rejected by device-id mismatch.
416 """
417 device_type, device_id = view.__dlpack_device__() 1lmnj
418 if device_type == _kDLCUDA and device_id != expected_device_id: 1almnj
419 raise ValueError( 1j
420 f"{operation} expects tensor on device {expected_device_id}, got {device_id}") 1j
421cdef inline intptr_t _get_current_context_ptr() except? 0:
422 cdef cydriver.CUcontext ctx
423 with nogil:
424 HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
425 if ctx == NULL:
426 raise RuntimeError("TensorMapDescriptor requires an active CUDA context")
427 return <intptr_t>ctx
430cdef inline int _get_current_device_id() except -1:
431 cdef cydriver.CUdevice dev
432 with nogil:
433 HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
434 return <int>dev
436def _compute_byte_strides(shape, strides, elem_size):
437 """Compute byte strides from element strides or C-contiguous fallback.
439 Returns a tuple of byte strides in row-major order.
440 """
441 if strides is not None:
442 return tuple(s * elem_size for s in strides)
444 # C-contiguous: compute byte strides from shape, innermost first
445 rank = len(shape)
446 byte_strides = []
447 stride = elem_size
448 for i in range(rank - 1, -1, -1):
449 byte_strides.append(stride)
450 stride *= shape[i]
451 byte_strides.reverse()
452 return tuple(byte_strides)
455def _validate_element_strides(element_strides, rank):
456 """Validate or default element_strides to all-ones."""
457 if element_strides is not None:
458 if len(element_strides) != rank:
459 raise ValueError(
460 f"element_strides must have {rank} elements, got {len(element_strides)}")
461 return element_strides
462 return (1,) * rank
465cdef class TensorMapDescriptor:
466 """Describes a TMA (Tensor Memory Accelerator) tensor map for Hopper+ GPUs.
468 A ``TensorMapDescriptor`` wraps the opaque 128-byte ``CUtensorMap`` struct
469 used by the hardware TMA unit for efficient bulk data movement between
470 global and shared memory.
472 Public tiled descriptors are created via
473 :meth:`cuda.core.StridedMemoryView.as_tensor_map`. Specialized
474 ``_from_*`` helpers remain private while this API surface settles, and
475 descriptors can be passed directly to :func:`~cuda.core.launch` as a
476 kernel argument.
477 """
479 def __init__(self):
480 raise RuntimeError( 1o
481 "TensorMapDescriptor cannot be instantiated directly. "
482 "Use StridedMemoryView.as_tensor_map() instead.")
484 cdef void* _get_data_ptr(self):
485 return <void*>&self._tensor_map
487 cdef int _check_context_compat(self) except -1:
488 cdef cydriver.CUcontext current_ctx
489 cdef cydriver.CUdevice current_dev
490 if self._context == 0 and self._device_id < 0:
491 return 0
492 with nogil:
493 HANDLE_RETURN(cydriver.cuCtxGetCurrent(¤t_ctx))
494 if current_ctx == NULL:
495 raise RuntimeError("TensorMapDescriptor requires an active CUDA context")
496 if self._context != 0 and <intptr_t>current_ctx != self._context:
497 raise RuntimeError(
498 "TensorMapDescriptor was created in a different CUDA context")
499 with nogil:
500 HANDLE_RETURN(cydriver.cuCtxGetDevice(¤t_dev))
501 cdef int current_dev_id = <int>current_dev
502 if self._device_id >= 0 and current_dev_id != self._device_id:
503 raise RuntimeError(
504 f"TensorMapDescriptor belongs to device {self._device_id}, "
505 f"but current device is {current_dev_id}")
506 return 0
508 @property
509 def device(self) -> Device | None:
510 """Return the :obj:`~cuda.core.Device` associated with this descriptor."""
511 if self._device_id >= 0:
512 from cuda.core._device import Device
513 return Device(self._device_id)
514 return None
516 @classmethod
517 def _from_tiled(cls, view, box_dim=None, *,
518 options=None,
519 element_strides=None,
520 data_type=None,
521 interleave=TensorMapInterleave.NONE,
522 swizzle=TensorMapSwizzle.NONE,
523 l2_promotion=TensorMapL2Promotion.NONE,
524 oob_fill=TensorMapOOBFill.NONE):
525 """Create a tiled TMA descriptor from a validated view.
527 Parameters
528 ----------
529 view : StridedMemoryView
530 A device-accessible view with a 16-byte-aligned pointer.
531 box_dim : tuple of int, optional
532 The size of each tile dimension (in elements). Must have the
533 same rank as the tensor and each value must be in [1, 256].
534 Specified in the same (row-major) order as the tensor shape.
535 Required unless ``options`` is provided.
536 options : TensorMapDescriptorOptions or mapping, optional
537 Bundled tiled-descriptor options. When provided, do not also pass
538 ``box_dim`` or the individual option kwargs.
539 element_strides : tuple of int, optional
540 Per-dimension element traversal strides. Default is all 1s.
541 Specified in the same (row-major) order as the tensor shape.
542 data_type : dtype-like or TensorMapDataType, optional
543 Explicit dtype override. If ``None``, inferred from the tensor's
544 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
545 accepted for compatibility.
546 interleave : TensorMapInterleave
547 Interleave layout. Default ``NONE``.
548 swizzle : TensorMapSwizzle
549 Swizzle mode. Default ``NONE``.
550 l2_promotion : TensorMapL2Promotion
551 L2 promotion mode. Default ``NONE``.
552 oob_fill : TensorMapOOBFill
553 Out-of-bounds fill mode. Default ``NONE``.
555 Returns
556 -------
557 TensorMapDescriptor
559 Raises
560 ------
561 ValueError
562 If the tensor rank is outside [1, 5], the pointer is not
563 16-byte aligned, or dimension/stride constraints are violated.
564 """
565 cdef TensorMapDescriptor desc = cls.__new__(cls) 1ab
567 opts = _coerce_tensor_map_descriptor_options( 1b
568 box_dim,
569 options,
570 element_strides=element_strides,
571 data_type=data_type,
572 interleave=interleave,
573 swizzle=swizzle,
574 l2_promotion=l2_promotion,
575 oob_fill=oob_fill, 1b
576 )
577 box_dim = opts.box_dim 1b
578 element_strides = opts.element_strides 1b
579 data_type = opts.data_type 1b
580 interleave = opts.interleave 1b
581 swizzle = opts.swizzle 1b
582 l2_promotion = opts.l2_promotion 1b
583 oob_fill = opts.oob_fill 1ab
585 _validate_tensor_map_view(view) 1b
586 # Keep both the original tensor object and the validated view alive.
587 # For DLPack exporters, the view may hold the owning capsule whose
588 # deleter can free the backing allocation when released.
589 desc._source_ref = view.exporting_obj
590 desc._view_ref = view
591 desc._context = _get_current_context_ptr()
592 desc._device_id = _get_current_device_id()
593 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_tiled")
595 tma_dt = _resolve_data_type(view, data_type)
596 cdef int c_data_type_int = tma_dt
597 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int
599 cdef intptr_t global_address = view.ptr
600 shape = view.shape
602 cdef int rank = len(shape)
603 if rank < 1 or rank > 5:
604 raise ValueError(
605 f"Tensor rank must be between 1 and 5, got {rank}")
607 if len(box_dim) != rank:
608 raise ValueError(
609 f"box_dim must have {rank} elements (same as tensor rank), "
610 f"got {len(box_dim)}")
612 for i, bd in enumerate(box_dim):
613 if bd < 1 or bd > 256:
614 raise ValueError(
615 f"box_dim[{i}] must be in [1, 256], got {bd}")
617 cdef bint elem_strides_provided = element_strides is not None
618 element_strides = _validate_element_strides(element_strides, rank)
620 # Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible.
621 # This avoids maintaining a second, independent validation/encoding implementation.
622 cdef uint8_t dl_code
623 cdef uint8_t dl_bits
624 cdef uint16_t dl_lanes
625 cdef int64_t c_shape[5]
626 cdef int64_t c_strides[5]
627 cdef int c_box_sizes[5]
628 cdef int c_elem_strides[5]
629 cdef const int64_t* c_strides_ptr
630 cdef const int* c_elem_strides_ptr
631 cdef char errbuf[512]
632 cdef int i_cccl
633 cdef int device_type
634 cdef int c_device_id
635 cdef int dl_device_type
636 cdef int dl_device_id
637 cdef int c_cccl_interleave_int
638 cdef int c_cccl_swizzle_int
639 cdef int c_cccl_l2_promotion_int
640 cdef int c_cccl_oob_fill_int
641 cdef int rc
642 if _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes):
643 c_strides_ptr = NULL
644 c_elem_strides_ptr = NULL
645 errbuf[0] = 0
647 for i_cccl in range(rank):
648 c_shape[i_cccl] = <int64_t>shape[i_cccl]
649 c_box_sizes[i_cccl] = <int>box_dim[i_cccl]
650 if elem_strides_provided:
651 c_elem_strides[i_cccl] = <int>element_strides[i_cccl]
653 if view.strides is not None:
654 for i_cccl in range(rank):
655 c_strides[i_cccl] = <int64_t>view.strides[i_cccl]
656 c_strides_ptr = &c_strides[0]
658 if elem_strides_provided:
659 c_elem_strides_ptr = &c_elem_strides[0]
661 dl_device_type, dl_device_id = view.__dlpack_device__()
662 device_type = dl_device_type
663 c_device_id = dl_device_id
664 c_cccl_interleave_int = int(interleave)
665 c_cccl_swizzle_int = int(swizzle)
666 c_cccl_l2_promotion_int = int(l2_promotion)
667 c_cccl_oob_fill_int = int(oob_fill)
669 with nogil:
670 rc = cuda_core_cccl_make_tma_descriptor_tiled(
671 <void*>&desc._tensor_map,
672 <void*>global_address,
673 device_type,
674 c_device_id,
675 rank,
676 &c_shape[0],
677 c_strides_ptr,
678 dl_code,
679 dl_bits,
680 dl_lanes,
681 &c_box_sizes[0],
682 c_elem_strides_ptr,
683 c_cccl_interleave_int,
684 c_cccl_swizzle_int,
685 c_cccl_l2_promotion_int,
686 c_cccl_oob_fill_int,
687 &errbuf[0],
688 <size_t>sizeof(errbuf),
689 )
691 if rc == 0:
692 desc._repr_info = {
693 "method": "tiled",
694 "rank": rank,
695 "data_type": TensorMapDataType(tma_dt),
696 "swizzle": swizzle,
697 }
698 return desc
700 msg = errbuf[:].split(b"\0", 1)[0].decode("utf-8", errors="replace")
701 # If CCCL isn't available at build time, fall back to the direct
702 # driver API path to preserve functionality on older toolchains.
703 if "not available at build time" not in msg:
704 raise ValueError(f"Failed to build TMA descriptor via CCCL: {msg}")
706 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
707 byte_strides = _compute_byte_strides(shape, view.strides, elem_size)
709 # Reverse dimensions for column-major cuTensorMap convention
710 # Python/DLPack: row-major (dim 0 = outermost)
711 # cuTensorMap: column-major (dim 0 = innermost)
712 cdef uint64_t[5] c_global_dim
713 cdef uint64_t[4] c_global_strides # rank - 1 elements
714 cdef uint32_t[5] c_box_dim
715 cdef uint32_t[5] c_element_strides
716 cdef int i_c
718 for i_c in range(rank):
719 # Reverse: Python dim i -> cuTensorMap dim (rank - 1 - i)
720 c_global_dim[i_c] = <uint64_t>shape[rank - 1 - i_c]
721 c_box_dim[i_c] = <uint32_t>box_dim[rank - 1 - i_c]
722 c_element_strides[i_c] = <uint32_t>element_strides[rank - 1 - i_c]
724 # globalStrides: rank-1 elements (byte strides for dims 1..N-1 in col-major order)
725 # The innermost stride (dim 0) is implicit = element size
726 for i_c in range(rank - 1):
727 c_global_strides[i_c] = <uint64_t>byte_strides[rank - 2 - i_c]
729 cdef uint32_t c_rank = <uint32_t>rank
730 cdef int c_interleave_int = int(interleave)
731 cdef int c_swizzle_int = int(swizzle)
732 cdef int c_l2_promotion_int = int(l2_promotion)
733 cdef int c_oob_fill_int = int(oob_fill)
734 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave>c_interleave_int
735 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle>c_swizzle_int
736 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion>c_l2_promotion_int
737 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill>c_oob_fill_int
739 with nogil:
740 HANDLE_RETURN(cydriver.cuTensorMapEncodeTiled(
741 &desc._tensor_map,
742 c_data_type,
743 c_rank,
744 <void*>global_address,
745 c_global_dim,
746 c_global_strides,
747 c_box_dim,
748 c_element_strides,
749 c_interleave,
750 c_swizzle,
751 c_l2_promotion,
752 c_oob_fill,
753 ))
755 desc._repr_info = {
756 "method": "tiled",
757 "rank": rank,
758 "data_type": TensorMapDataType(tma_dt),
759 "swizzle": swizzle,
760 }
762 return desc
764 @classmethod
765 def _from_im2col(cls, view, pixel_box_lower_corner, pixel_box_upper_corner,
766 channels_per_pixel, pixels_per_column, *,
767 element_strides=None,
768 data_type=None,
769 interleave=TensorMapInterleave.NONE,
770 swizzle=TensorMapSwizzle.NONE,
771 l2_promotion=TensorMapL2Promotion.NONE,
772 oob_fill=TensorMapOOBFill.NONE):
773 """Create an im2col TMA descriptor from a validated view.
775 Im2col layout is used for convolution-style data access patterns.
777 Parameters
778 ----------
779 view : StridedMemoryView
780 A device-accessible view with a 16-byte-aligned pointer.
781 pixel_box_lower_corner : tuple of int
782 Lower corner of the pixel bounding box for each spatial
783 dimension (rank - 2 elements). Specified in row-major order
784 matching the tensor's spatial dimensions.
785 pixel_box_upper_corner : tuple of int
786 Upper corner of the pixel bounding box for each spatial
787 dimension (rank - 2 elements). Specified in row-major order
788 matching the tensor's spatial dimensions.
789 channels_per_pixel : int
790 Number of channels per pixel.
791 pixels_per_column : int
792 Number of pixels per column.
793 element_strides : tuple of int, optional
794 Per-dimension element traversal strides. Default is all 1s.
795 data_type : dtype-like or TensorMapDataType, optional
796 Explicit dtype override. If ``None``, inferred from the tensor's
797 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
798 accepted for compatibility.
799 interleave : TensorMapInterleave
800 Interleave layout. Default ``NONE``.
801 swizzle : TensorMapSwizzle
802 Swizzle mode. Default ``NONE``.
803 l2_promotion : TensorMapL2Promotion
804 L2 promotion mode. Default ``NONE``.
805 oob_fill : TensorMapOOBFill
806 Out-of-bounds fill mode. Default ``NONE``.
808 Returns
809 -------
810 TensorMapDescriptor
812 Raises
813 ------
814 ValueError
815 If the tensor rank is outside [3, 5], the pointer is not
816 16-byte aligned, or other constraints are violated.
817 """
818 cdef TensorMapDescriptor desc = cls.__new__(cls)
820 _validate_tensor_map_view(view)
821 desc._source_ref = view.exporting_obj
822 desc._view_ref = view
823 desc._context = _get_current_context_ptr()
824 desc._device_id = _get_current_device_id()
825 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_im2col")
827 tma_dt = _resolve_data_type(view, data_type)
828 cdef int c_data_type_int = tma_dt
829 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int
831 cdef intptr_t global_address = view.ptr
832 shape = view.shape
834 cdef int rank = len(shape)
835 if rank < 3 or rank > 5:
836 raise ValueError(
837 f"Im2col tensor rank must be between 3 and 5, got {rank}")
839 cdef int n_spatial = rank - 2
840 if len(pixel_box_lower_corner) != n_spatial:
841 raise ValueError(
842 f"pixel_box_lower_corner must have {n_spatial} elements "
843 f"(rank - 2), got {len(pixel_box_lower_corner)}")
844 if len(pixel_box_upper_corner) != n_spatial:
845 raise ValueError(
846 f"pixel_box_upper_corner must have {n_spatial} elements "
847 f"(rank - 2), got {len(pixel_box_upper_corner)}")
849 element_strides = _validate_element_strides(element_strides, rank)
851 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
852 byte_strides = _compute_byte_strides(shape, view.strides, elem_size)
854 # Reverse all dimension arrays for column-major convention
855 cdef uint64_t[5] c_global_dim
856 cdef uint64_t[4] c_global_strides
857 cdef uint32_t[5] c_element_strides
858 cdef int[3] c_pixel_box_lower # max 3 spatial dims (rank 5 - 2)
859 cdef int[3] c_pixel_box_upper
860 cdef int i_c
862 for i_c in range(3):
863 c_pixel_box_lower[i_c] = 0
864 c_pixel_box_upper[i_c] = 0
866 for i_c in range(rank):
867 c_global_dim[i_c] = <uint64_t>shape[rank - 1 - i_c]
868 c_element_strides[i_c] = <uint32_t>element_strides[rank - 1 - i_c]
870 for i_c in range(rank - 1):
871 c_global_strides[i_c] = <uint64_t>byte_strides[rank - 2 - i_c]
873 # Reverse spatial dimensions for lower/upper corners
874 for i_c in range(n_spatial):
875 c_pixel_box_lower[i_c] = <int>pixel_box_lower_corner[n_spatial - 1 - i_c]
876 c_pixel_box_upper[i_c] = <int>pixel_box_upper_corner[n_spatial - 1 - i_c]
878 cdef uint32_t c_rank = <uint32_t>rank
879 cdef uint32_t c_channels = <uint32_t>channels_per_pixel
880 cdef uint32_t c_pixels = <uint32_t>pixels_per_column
881 cdef int c_interleave_int = int(interleave)
882 cdef int c_swizzle_int = int(swizzle)
883 cdef int c_l2_promotion_int = int(l2_promotion)
884 cdef int c_oob_fill_int = int(oob_fill)
885 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave>c_interleave_int
886 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle>c_swizzle_int
887 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion>c_l2_promotion_int
888 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill>c_oob_fill_int
890 with nogil:
891 HANDLE_RETURN(cydriver.cuTensorMapEncodeIm2col(
892 &desc._tensor_map,
893 c_data_type,
894 c_rank,
895 <void*>global_address,
896 c_global_dim,
897 c_global_strides,
898 c_pixel_box_lower,
899 c_pixel_box_upper,
900 c_channels,
901 c_pixels,
902 c_element_strides,
903 c_interleave,
904 c_swizzle,
905 c_l2_promotion,
906 c_oob_fill,
907 ))
909 desc._repr_info = {
910 "method": "im2col",
911 "rank": rank,
912 "data_type": TensorMapDataType(tma_dt),
913 "swizzle": swizzle,
914 }
916 return desc
918 @classmethod
919 def _from_im2col_wide(cls, view, pixel_box_lower_corner_width, pixel_box_upper_corner_width,
920 channels_per_pixel, pixels_per_column, *,
921 element_strides=None,
922 data_type=None,
923 interleave=TensorMapInterleave.NONE,
924 mode=TensorMapIm2ColWideMode.W,
925 swizzle=TensorMapSwizzle.SWIZZLE_128B,
926 l2_promotion=TensorMapL2Promotion.NONE,
927 oob_fill=TensorMapOOBFill.NONE):
928 """Create an im2col-wide TMA descriptor from a validated view.
930 Im2col-wide layout loads elements exclusively along the W (width)
931 dimension. This variant is supported on compute capability 10.0+
932 (Blackwell and later).
934 Parameters
935 ----------
936 view : StridedMemoryView
937 A device-accessible view with a 16-byte-aligned pointer.
938 pixel_box_lower_corner_width : int
939 Lower corner of the pixel bounding box along the W dimension.
940 pixel_box_upper_corner_width : int
941 Upper corner of the pixel bounding box along the W dimension.
942 channels_per_pixel : int
943 Number of channels per pixel.
944 pixels_per_column : int
945 Number of pixels per column.
946 element_strides : tuple of int, optional
947 Per-dimension element traversal strides. Default is all 1s.
948 data_type : dtype-like or TensorMapDataType, optional
949 Explicit dtype override. If ``None``, inferred from the tensor's
950 dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
951 accepted for compatibility.
952 interleave : TensorMapInterleave
953 Interleave layout. Default ``NONE``.
954 mode : TensorMapIm2ColWideMode
955 Im2col wide mode. Default ``W``.
956 swizzle : TensorMapSwizzle
957 Swizzle mode. Default ``SWIZZLE_128B``.
958 l2_promotion : TensorMapL2Promotion
959 L2 promotion mode. Default ``NONE``.
960 oob_fill : TensorMapOOBFill
961 Out-of-bounds fill mode. Default ``NONE``.
963 Returns
964 -------
965 TensorMapDescriptor
967 Raises
968 ------
969 ValueError
970 If the tensor rank is outside [3, 5], the pointer is not
971 16-byte aligned, or other constraints are violated.
972 """
973 IF CUDA_CORE_BUILD_MAJOR < 13:
974 raise RuntimeError(
975 "TensorMapDescriptor._from_im2col_wide requires a CUDA 13+ build")
976 ELSE:
977 cdef TensorMapDescriptor desc = cls.__new__(cls)
979 _validate_tensor_map_view(view)
980 desc._source_ref = view.exporting_obj
981 desc._view_ref = view
982 desc._context = _get_current_context_ptr()
983 desc._device_id = _get_current_device_id()
984 _require_view_device(view, desc._device_id, "TensorMapDescriptor._from_im2col_wide")
986 tma_dt = _resolve_data_type(view, data_type)
987 cdef int c_data_type_int = tma_dt
988 cdef cydriver.CUtensorMapDataType c_data_type = <cydriver.CUtensorMapDataType>c_data_type_int
990 cdef intptr_t global_address = view.ptr
991 shape = view.shape
993 cdef int rank = len(shape)
994 if rank < 3 or rank > 5:
995 raise ValueError(
996 f"Im2col-wide tensor rank must be between 3 and 5, got {rank}")
998 element_strides = _validate_element_strides(element_strides, rank)
1000 cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
1001 byte_strides = _compute_byte_strides(shape, view.strides, elem_size)
1003 # Reverse all dimension arrays for column-major convention
1004 cdef uint64_t[5] c_global_dim
1005 cdef uint64_t[4] c_global_strides
1006 cdef uint32_t[5] c_element_strides
1007 cdef int i_c
1009 for i_c in range(rank):
1010 c_global_dim[i_c] = <uint64_t>shape[rank - 1 - i_c]
1011 c_element_strides[i_c] = <uint32_t>element_strides[rank - 1 - i_c]
1013 for i_c in range(rank - 1):
1014 c_global_strides[i_c] = <uint64_t>byte_strides[rank - 2 - i_c]
1016 cdef uint32_t c_rank = <uint32_t>rank
1017 cdef int c_lower_w = <int>pixel_box_lower_corner_width
1018 cdef int c_upper_w = <int>pixel_box_upper_corner_width
1019 cdef uint32_t c_channels = <uint32_t>channels_per_pixel
1020 cdef uint32_t c_pixels = <uint32_t>pixels_per_column
1021 cdef int c_interleave_int = int(interleave)
1022 cdef int c_mode_int = int(mode)
1023 cdef int c_swizzle_int = int(swizzle)
1024 cdef int c_l2_promotion_int = int(l2_promotion)
1025 cdef int c_oob_fill_int = int(oob_fill)
1026 cdef cydriver.CUtensorMapInterleave c_interleave = <cydriver.CUtensorMapInterleave>c_interleave_int
1027 cdef cydriver.CUtensorMapIm2ColWideMode c_mode = <cydriver.CUtensorMapIm2ColWideMode>c_mode_int
1028 cdef cydriver.CUtensorMapSwizzle c_swizzle = <cydriver.CUtensorMapSwizzle>c_swizzle_int
1029 cdef cydriver.CUtensorMapL2promotion c_l2_promotion = <cydriver.CUtensorMapL2promotion>c_l2_promotion_int
1030 cdef cydriver.CUtensorMapFloatOOBfill c_oob_fill = <cydriver.CUtensorMapFloatOOBfill>c_oob_fill_int
1032 with nogil:
1033 HANDLE_RETURN(cydriver.cuTensorMapEncodeIm2colWide(
1034 &desc._tensor_map,
1035 c_data_type,
1036 c_rank,
1037 <void*>global_address,
1038 c_global_dim,
1039 c_global_strides,
1040 c_lower_w,
1041 c_upper_w,
1042 c_channels,
1043 c_pixels,
1044 c_element_strides,
1045 c_interleave,
1046 c_mode,
1047 c_swizzle,
1048 c_l2_promotion,
1049 c_oob_fill,
1050 ))
1052 desc._repr_info = {
1053 "method": "im2col_wide",
1054 "rank": rank,
1055 "data_type": TensorMapDataType(tma_dt),
1056 "swizzle": swizzle,
1057 }
1059 return desc
1061 def replace_address(self, tensor: object) -> None:
1062 """Replace the global memory address in this tensor map descriptor.
1064 This is useful when the tensor data has been reallocated but the
1065 shape, strides, and other parameters remain the same.
1067 Parameters
1068 ----------
1069 tensor : object
1070 Any object supporting DLPack or ``__cuda_array_interface__``,
1071 or a :obj:`~cuda.core.StridedMemoryView`. Must refer to
1072 device-accessible memory with a 16-byte-aligned pointer.
1073 """
1074 self._check_context_compat()
1075 view = _get_validated_view(tensor)
1076 _require_view_device(view, self._device_id, "replace_address")
1078 cdef intptr_t global_address = view.ptr
1080 with nogil:
1081 HANDLE_RETURN(cydriver.cuTensorMapReplaceAddress(
1082 &self._tensor_map,
1083 <void*>global_address,
1084 ))
1086 # Update the source reference only after the driver call succeeds,
1087 # so we don't drop the old tensor (risking a dangling pointer in the
1088 # CUtensorMap struct) if the call fails.
1089 self._source_ref = view.exporting_obj
1090 self._view_ref = view
1092 def __repr__(self) -> str:
1093 info = self._repr_info
1094 if info is None:
1095 return "TensorMapDescriptor()"
1096 parts = []
1097 if "method" in info:
1098 parts.append(info["method"])
1099 if "rank" in info:
1100 parts.append(f"rank={info['rank']}")
1101 if "data_type" in info:
1102 parts.append(f"dtype={info['data_type'].name}")
1103 if "swizzle" in info:
1104 parts.append(f"swizzle={info['swizzle'].name}")
1105 return f"TensorMapDescriptor({', '.join(parts)})"