Mojo function
create_tma_tile_template
create_tma_tile_template[dtype: DType, rank: Int, tile_shape: IndexList[rank], /, 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]()]() -> TMATensorTile[dtype, rank, __tile_shape, __desc_shape]
Same as create_tma_tile expect the descriptor is only a placeholder or a template for later replacement.
specification of data type, rank, and layout orientation. It supports both 2D and 3D tensors and provides fine-grained control over the memory access patterns.
Constraints:
- Only supports 2D and 3D tensors (rank must be 2 or 3).
- 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 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 or 3). - tile_shape (
IndexList): IndexList[rank] The shape of the tile to be transferred. - swizzle_mode (
TensorMapSwizzle): TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE The swizzling mode to use for memory access optimization. - __tile_shape (
IndexList): IndexList[rank] = tile_shape Internal parameter for the tile shape in shared memory. - __desc_shape (
IndexList): IndexList[rank] = _default_desc_shape... Internal parameter for the descriptor shape, which may differ from the tile shape to accommodate hardware requirements.
Returns:
TMATensorTile: A TMATensorTile configured with the specified parameters, ready for use in
asynchronous data transfer operations.
Raises:
If TMA descriptor creation fails.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!