Skip to main content

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.

Was this page helpful?