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_tensor_tile

def create_tensor_tile[dtype: DType, rank: Int, //, tile_shape: IndexList[rank], /, k_major_tma: Bool = True, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, *, __tile_shape: IndexList[rank] = tile_shape, __desc_shape: IndexList[rank] = _default_desc_shape[rank, dtype, tile_shape, swizzle_mode]()](ctx: DeviceContext, tensor: LayoutTensor[dtype, address_space=tensor.address_space, element_layout=tensor.element_layout, layout_int_type=tensor.layout_int_type, linear_idx_type=tensor.linear_idx_type, masked=tensor.masked, alignment=tensor.alignment]) -> TMATensorTile[dtype, rank, __tile_shape, __desc_shape, k_major_tma]

Creates a TMATensorTile with advanced configuration options for 2D, 3D, 4D, or 5D tensors.

This overload provides more control over the TMA descriptor creation, allowing specification of data type, rank, and layout orientation. It supports 2D, 3D, 4D, and 5D tensors and provides fine-grained control over the memory access patterns.

Constraints:

• Only supports 2D, 3D, 4D, and 5D tensors (rank must be 2, 3, 4, or 5).
• For non-SWIZZLE_NONE modes, the K dimension size in bytes must be a multiple of the swizzle mode's byte size.
• For MN-major layout, only SWIZZLE_128B is supported.
• For 3D, 4D, and 5D tensors, only K-major layout is supported.

Parameters:

• ​dtype (DType): DType The data type of the tensor elements.
• ​rank (Int): Int The dimensionality of the tensor (must be 2, 3, 4, or 5).
• ​tile_shape (IndexList[rank]): IndexList[rank] The shape of the tile to be transferred.
• ​k_major_tma (Bool): Bool = True Whether the tma should copy desc into shared memory following a column-major (if True) or row-major (if False) pattern.
• ​swizzle_mode (TensorMapSwizzle): TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE The swizzling mode to use for memory access optimization.
• ​__tile_shape (IndexList[rank]): IndexList[rank] = tile_shape Internal parameter for the tile shape in shared memory.
• ​__desc_shape (IndexList[rank]): IndexList[rank] = _default_desc_shape... Internal parameter for the descriptor shape, which may differ from the tile shape to accommodate hardware requirements.

Args:

• ​ctx (DeviceContext): DeviceContext The CUDA device context used to create the TMA descriptor.
• ​tensor (LayoutTensor[dtype, address_space=tensor.address_space, element_layout=tensor.element_layout, layout_int_type=tensor.layout_int_type, linear_idx_type=tensor.linear_idx_type, masked=tensor.masked, alignment=tensor.alignment]): LayoutTensor[dtype, ...] The source tensor from which data will be transferred. This defines the global memory layout and must match the specified data type.

Returns:

TMATensorTile[dtype, rank, __tile_shape, __desc_shape, k_major_tma]: A TMATensorTile configured with the specified parameters, ready for use in asynchronous data transfer operations.

Raises:

If TMA descriptor creation fails.

def create_tensor_tile[dtype: DType, rank: Int, //, tile_shape: IndexList[rank], /, k_major_tma: Bool = True, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, *, __tile_shape: IndexList[rank] = tile_shape, __desc_shape: IndexList[rank] = _default_desc_shape[rank, dtype, tile_shape, swizzle_mode]()](ctx: DeviceContext, tensor: TileTensor[dtype, Storage=tensor.Storage, address_space=tensor.address_space, linear_idx_type=tensor.linear_idx_type]) -> TMATensorTile[dtype, rank, __tile_shape, __desc_shape, k_major_tma]

Creates a TMATensorTile from a TileTensor.

This overload accepts a TileTensor instead of LayoutTensor, enabling use with the new coordinate-based tensor abstraction.

Parameters:

• ​dtype (DType): The data type of the tensor elements.
• ​rank (Int): The dimensionality of the tensor (must be 2, 3, 4, or 5).
• ​tile_shape (IndexList[rank]): The shape of the tile to be transferred.
• ​k_major_tma (Bool): Whether the TMA should use column-major pattern.
• ​swizzle_mode (TensorMapSwizzle): The swizzling mode for memory access optimization.
• ​__tile_shape (IndexList[rank]): Internal parameter for the tile shape.
• ​__desc_shape (IndexList[rank]): Internal parameter for the descriptor shape.

Args:

• ​ctx (DeviceContext): The CUDA device context.
• ​tensor (TileTensor[dtype, Storage=tensor.Storage, address_space=tensor.address_space, linear_idx_type=tensor.linear_idx_type]): The source TileTensor.

Returns:

TMATensorTile[dtype, rank, __tile_shape, __desc_shape, k_major_tma]: A TMATensorTile configured for the given tensor.

Raises:

If TMA descriptor creation fails.