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 function

consumer_main_loop

def consumer_main_loop[accum_type: DType, c_type: DType, a_type: DType, b_type: DType, a_dim0: Int, a_dim1: Int, a_num_tiles: Int, a_swizzle_bytes: Int, b_dim0: Int, b_dim1: Int, b_num_tiles: Int, b_swizzle_bytes: Int, a_swizzle: TensorMapSwizzle, b_swizzle: TensorMapSwizzle, transpose_b: Bool, pipeline_stages: Int, //, *, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], cta_group: Int = Int(1), cluster_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(1), Int(1), Int(1))](tmem_addr: UInt32, a_smem_tiles: SMemTileArray2D[a_type, a_dim0, a_dim1, a_num_tiles, a_swizzle_bytes], b_smem_tiles: SMemTileArray2D[b_type, b_dim0, b_dim1, b_num_tiles, b_swizzle_bytes], mma_mbar: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED, _safe=False], tma_mbar: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED, _safe=False], consumer_phase: PipelineState[pipeline_stages], mma_op: MmaOpSM100_SS[c_type, a_type, b_type, block_tile_shape, mma_shape, accum_type=accum_type, cta_group=cta_group, cluster_shape=cluster_shape, a_swizzle=a_swizzle, b_swizzle=b_swizzle, transpose_b=transpose_b], elect_one_warp: Bool, iter_idx: UInt32)

Performs the MMA consumer step for one pipeline stage.

Waits on the TMA mbarrier for the current stage, issues the SM100 SS MMA against the staged shared-memory tiles into tensor memory, and commits the result to the MMA mbarrier.

Parameters:

accum_type (DType): Accumulator dtype used for the MMA result in tensor memory (inferred).
c_type (DType): Element type of the C output matrix (inferred).
a_type (DType): Element type of the A operand tiles (inferred).
b_type (DType): Element type of the B operand tiles (inferred).
a_dim0 (Int): Number of rows in one A shared-memory tile (inferred).
a_dim1 (Int): Number of columns in one A shared-memory tile (inferred).
a_num_tiles (Int): Number of pipeline stages in the A shared-memory tile array (inferred).
a_swizzle_bytes (Int): Swizzle granularity in bytes for A shared-memory tiles (inferred).
b_dim0 (Int): Number of rows in one B shared-memory tile (inferred).
b_dim1 (Int): Number of columns in one B shared-memory tile (inferred).
b_num_tiles (Int): Number of pipeline stages in the B shared-memory tile array (inferred).
b_swizzle_bytes (Int): Swizzle granularity in bytes for B shared-memory tiles (inferred).
a_swizzle (TensorMapSwizzle): TMA swizzle mode applied to A shared-memory tiles (inferred).
b_swizzle (TensorMapSwizzle): TMA swizzle mode applied to B shared-memory tiles (inferred).
transpose_b (Bool): Whether B is stored transposed in global memory (inferred).
pipeline_stages (Int): Number of double-buffered pipeline stages for A and B shared-memory tiles (inferred).
block_tile_shape (IndexList[Int(3)]): Block tile shape [BM, BN, BK] partitioning the GEMM into work tiles.
mma_shape (IndexList[Int(3)]): MMA instruction shape [MMA_M, MMA_N, MMA_K].
cta_group (Int): Number of CTAs cooperating per MMA along the M dimension (defaults to 1).
cluster_shape (IndexList[Int(3)]): Thread block cluster shape [CLUSTER_M, CLUSTER_N, CLUSTER_D] (defaults to Index(1, 1, 1)).

Args:

tmem_addr (UInt32): Tensor-memory base address where the MMA accumulates results for this accumulator stage.
a_smem_tiles (SMemTileArray2D[a_type, a_dim0, a_dim1, a_num_tiles, a_swizzle_bytes]): Shared-memory tile array of staged A tiles, one per pipeline stage.
b_smem_tiles (SMemTileArray2D[b_type, b_dim0, b_dim1, b_num_tiles, b_swizzle_bytes]): Shared-memory tile array of staged B tiles, one per pipeline stage.
mma_mbar (Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED, _safe=False]): Pointer to the MMA mbarrier array, one per pipeline stage, signaled when the MMA for a stage completes.
tma_mbar (Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED, _safe=False]): Pointer to the TMA mbarrier array, one per pipeline stage, waited on to confirm the TMA load for a stage is complete.
consumer_phase (PipelineState[pipeline_stages]): Consumer pipeline state tracking the stage index and phase bit.
mma_op (MmaOpSM100_SS[c_type, a_type, b_type, block_tile_shape, mma_shape, accum_type=accum_type, cta_group=cta_group, cluster_shape=cluster_shape, a_swizzle=a_swizzle, b_swizzle=b_swizzle, transpose_b=transpose_b]): SM100 SS MMA operation object used to issue and commit the tensor-memory multiply-accumulate.
elect_one_warp (Bool): Whether this warp is the elected MMA warp for the current work tile.
iter_idx (UInt32): Current K-dimension iteration index; when zero, the MMA initializes the accumulator instead of accumulating.