Skip to main content

Mojo function

consumer_main_loop

consumer_main_loop[accum_type: DType, c_type: DType, a_type: DType, b_type: DType, sfa_dtype: DType, sfb_dtype: 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, scaling_kind: UMMAKind, /, *, block_tile_shape: IndexList[3], mma_shape: IndexList[3], SFA_NUM_COLS: Int, SFB_NUM_COLS: Int, cta_group: Int = 1, cluster_shape: IndexList[3] = Index(1, 1, 1), k_group_size: UInt = 1](tmem_addr: UInt32, sfa_tmem: UInt32, sfb_tmem: 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], sfa_smem_tiles: SMemTileArrayWithLayout[sfa_dtype, sfa_smem_tiles.tile_layout, sfa_smem_tiles.num_tiles, sfa_smem_tiles.alignment], sfb_smem_tiles: SMemTileArrayWithLayout[sfb_dtype, sfb_smem_tiles.tile_layout, sfb_smem_tiles.num_tiles, sfb_smem_tiles.alignment], load_mma_pipeline: ProducerConsumerPipeline[pipeline_stages], mma_op: MmaOpSM100_BlockScaled_SS[c_type, a_type, b_type, sfa_dtype, sfb_dtype, scaling_kind, 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, k_start: UInt32, work_tile_coord: Tuple[UInt, UInt])

TileTensor-based consumer_main_loop for block-scaled MMA.

Accepts SMemTileArray2D for A/B tiles and SMemTileArrayWithLayout for scale factor tiles, calling the TileTensor MMA path.

Was this page helpful?