Mojo function
consumer_main_loop
consumer_main_loop[accum_type: DType, c_type: DType, a_type: DType, b_type: DType, a_smem_layout: Layout, b_smem_layout: Layout, a_swizzle: TensorMapSwizzle, b_swizzle: TensorMapSwizzle, transpose_b: Bool, pipeline_stages: Int, /, *, block_tile_shape: IndexList[3], mma_shape: IndexList[3], cta_group: Int = 1, cluster_shape: IndexList[3] = Index(1, 1, 1), k_group_size: UInt = 1](tmem_addr: UInt32, a_smem_iter: LayoutTensorIter[a_type, a_smem_layout, MutAnyOrigin, address_space=AddressSpace.SHARED, alignment=128], b_smem_iter: LayoutTensorIter[b_type, b_smem_layout, MutAnyOrigin, address_space=AddressSpace.SHARED, alignment=128], load_mma_pipeline: ProducerConsumerPipeline[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, k_start: UInt32)
Consume tiles from shared memory and execute MMA operations.
This is the public API for external callers using SMemTileIter.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!