IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

MatmulTileWriter

struct MatmulTileWriter[dtype: DType, tensor_layout: TensorLayout, tensor_storage: TensorStorage, linear_idx_type: DType, smem_tile_layout: TensorLayout, //, *, BM: Int, BN: Int, swizzle: TensorMapSwizzle, wgmma_shape: IndexList[Int(3)], num_consumer: Int = Int(1), use_tma_store: Bool = False, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]] = None, swapAB: Bool = False]

Writes WGMMA accumulator register tiles to global memory for SM90 matmul.

Provides two write paths: TMA store (shared memory → global via hardware async copies) and thread-wise store (register → global directly). Handles swizzled shared memory layouts, optional elementwise epilogues, and A/B swap when the problem uses the swapAB small-M strategy.

Parameters​

  • ​dtype (DType): Element type of the output tensor and shared memory tile (inferred).
  • ​tensor_layout (TensorLayout): Memory layout of the output tensor in global memory (inferred).
  • ​tensor_storage (TensorStorage): Storage backing the output tensor in global memory (inferred).
  • ​linear_idx_type (DType): Integer type used for linear index arithmetic into the output tensor (inferred).
  • ​smem_tile_layout (TensorLayout): Memory layout of the shared memory tile used to stage output before storing to global memory; also defines the workgroup tile dimensions WG_BM and WG_BN (inferred).
  • ​BM (Int): Row (M) dimension of the matmul output tile per block.
  • ​BN (Int): Column (N) dimension of the matmul output tile per block.
  • ​swizzle (TensorMapSwizzle): TMA swizzle mode applied to shared memory for TMA async stores.
  • ​wgmma_shape (IndexList[Int(3)]): Shape of each WGMMA instruction as a 3-element index list (M, N, K).
  • ​num_consumer (Int): Number of consumer warp groups sharing the output tile (defaults to 1).
  • ​use_tma_store (Bool): Whether to use TMA async stores for shared memory to global memory copies (defaults to False).
  • ​elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue lambda invoked with each output element's global coordinates and value, writing directly to global memory (defaults to None).
  • ​elementwise_compute_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]]): Optional compute lambda that transforms each output element value before it is written back to shared memory and then to global memory (defaults to None).
  • ​swapAB (Bool): Whether to swap the A and B operand roles for the small-M strategy, transposing the tile and block coordinate mapping (defaults to False).

Fields​

  • ​tensor (MatmulTileWriter[BM=BM, BN=BN, swizzle=swizzle, wgmma_shape=wgmma_shape, num_consumer=num_consumer, use_tma_store=use_tma_store, elementwise_lambda_fn=elementwise_lambda_fn, elementwise_compute_lambda_fn=elementwise_compute_lambda_fn, swapAB=swapAB].CTensorType):
  • ​smem_tile (TileTensor[dtype, smem_tile_layout, MutAnyOrigin, address_space=AddressSpace.SHARED]):
  • ​warp_group_thread_idx (Int):
  • ​local_warp_group_idx (Int):
  • ​local_thread_idx (Int):
  • ​block_y (Int):
  • ​block_x (Int):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

CTensorType​

comptime CTensorType = TileTensor[dtype, tensor_layout, MutAnyOrigin, Storage=tensor_storage, linear_idx_type=linear_idx_type]

frag_size​

comptime frag_size = (Int((mul wgmma_shape[Int(0)], wgmma_shape[Int(1)])) // _resolve_warpgroup_size())

lambda_type​

comptime lambda_type = def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], mut SIMD[dtype, width]) -> None

N​

comptime N = tensor_layout.static_shape[Int(1)]

num_consumer_threads​

comptime num_consumer_threads = (num_consumer * _resolve_warpgroup_size())

num_m_mmas​

comptime num_m_mmas = ((BM // wgmma_shape[Int(0)]) // num_consumer)

num_n_mmas​

comptime num_n_mmas = (BN // wgmma_shape[Int(1)])

simd_size​

comptime simd_size = simd_width_of[dtype]()

WG_BM​

comptime WG_BM = smem_tile_layout.static_shape[Int(0)]

WG_BN​

comptime WG_BN = smem_tile_layout.static_shape[Int(1)]

Methods​

__init__​

def __init__(tensor: TileTensor[dtype, tensor_layout, MutAnyOrigin, Storage=tensor_storage, linear_idx_type=linear_idx_type], smem_tile: TileTensor[dtype, smem_tile_layout, MutAnyOrigin, address_space=AddressSpace.SHARED], warp_group_thread_idx: Int, local_warp_group_idx: Int, local_thread_idx: Int, block_y: Int, block_x: Int) -> Self

write_tile​

def write_tile[tma_rank: Int, tma_tile_shape: IndexList[tma_rank], tma_desc_shape: IndexList[tma_rank], accum_type: DType, reg_tile_layout: Layout, //](self, tma_op: TMATensorTile[dtype, tma_rank, tma_tile_shape, tma_desc_shape], reg_tile: LayoutTensor[accum_type, reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL])

Write output from registers to global memory.

Selects optimized st.matrix path for bf16 when constraints are met, otherwise uses general register-to-global path.

Parameters:

  • ​tma_rank (Int): Number of dimensions in the TMA tensor descriptor.
  • ​tma_tile_shape (IndexList[tma_rank]): Shape of each TMA store tile per async copy, as an index list of length tma_rank.
  • ​tma_desc_shape (IndexList[tma_rank]): Full shape of the TMA tensor descriptor as an index list of length tma_rank.
  • ​accum_type (DType): Data type of the WGMMA accumulator register tile.
  • ​reg_tile_layout (Layout): Memory layout of the accumulator register tile.

Args:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact