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)](tmem_addr: SIMD[uint32, 1], a_smem_iter: LayoutTensorIter[a_type, a_smem_layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128], b_smem_iter: LayoutTensorIter[b_type, b_smem_layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128], mma_mbar: UnsafePointer[SharedMemBarrier, address_space=AddressSpace(3), alignment=16], tma_mbar: UnsafePointer[SharedMemBarrier, address_space=AddressSpace(3), alignment=16], 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: UInt, accum_index: UInt, stage_stride_cols: UInt)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!