Skip to main content
Log in

Mojo function

consumer_main_loop

consumer_main_loop[accum_type: DType, a_type: DType, b_type: DType, c_reg_layout: Layout, a_smem_layout: Layout, b_smem_layout: Layout, wgmma_shape: Index[3], a_swizzle: TensorMapSwizzle, b_swizzle: TensorMapSwizzle, transpose_b: Bool, pipeline_stages: Int, /, *, cluster_shape: StaticTuple[SIMD[int32, 1], 3] = StaticTuple(__init__[__mlir_type.!pop.int_literal](1), __init__[__mlir_type.!pop.int_literal](1), __init__[__mlir_type.!pop.int_literal](1)), promotion_frequency: Int = 1, num_consumer: Int = 1](final_c_reg_tile: LayoutTensor[accum_type, c_reg_layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], c_reg_tile: LayoutTensor[accum_type, c_reg_layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], a_smem_iter: LayoutTensorIter[a_type, a_smem_layout, origin, address_space=AddressSpace(3), alignment=128, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked], b_smem_iter: LayoutTensorIter[b_type, b_smem_layout, origin, address_space=AddressSpace(3), alignment=128, circular=circular, axis=axis, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked], mut read_pipeline_states: PipelineState[pipeline_stages], full: UnsafePointer[SharedMemBarrier, address_space=AddressSpace(3), alignment=8], empty: UnsafePointer[SharedMemBarrier, address_space=AddressSpace(3), alignment=8], wgmma_op: TensorCoreAsync[accum_type, a_type, b_type, wgmma_shape, a_swizzle, b_swizzle, transpose_b], num_k_iters: Int, local_warp_group_idx: UInt, warp_group_thread_idx: UInt)