For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
create_tma_descriptor
def create_tma_descriptor[dtype: DType, rank: Int, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE](global_buf: DeviceBuffer[dtype], global_shape: IndexList[rank], global_strides: IndexList[rank], shared_mem_shape: IndexList[rank]) -> TMADescriptor
Creates a TMA descriptor for tiled memory operations.
Encodes tensor layout information into a 128-byte TMA descriptor that can be used with TMA hardware instructions to efficiently copy data between global and shared memory on NVIDIA GPUs.
The descriptor specifies a mapping from a tile in shared memory to a region in global memory, including dimensions, strides, data type, and optional swizzling for bank conflict avoidance.
Parameters:
- dtype (
DType): The element data type of the tensor. - rank (
Int): The number of dimensions (1-5). - swizzle_mode (
TensorMapSwizzle): The swizzle pattern to apply in shared memory. - l2_promotion (
TensorMapL2Promotion): L2 cache promotion hint for TMA loads. Defaults to NONE.
Args:
- global_buf (
DeviceBuffer[dtype]): Device buffer containing the global memory tensor. - global_shape (
IndexList[rank]): Dimensions of the tensor in global memory. - global_strides (
IndexList[rank]): Strides (in elements) for each dimension in global memory. The tensor must be row-major (stride at innermost dimension equals 1). - shared_mem_shape (
IndexList[rank]): Dimensions of the tile to be copied to shared memory.
Returns:
TMADescriptor: A TMA descriptor configured for the specified tensor layout.
Raises:
An error if the descriptor creation fails.