Mojo function
consumer_main_loop
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[3], mma_shape: IndexList[3], cta_group: Int = 1, cluster_shape: IndexList[3] = Index[Int, Int, Int](1, 1, 1), k_group_size: 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], 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)
TileTensor overload of consumer_main_loop.
Accepts SMemTileArray2D instead of LayoutTensorIter, indexing directly
into the tile arrays to get TileTensor tiles for MMA. The tile dimension
and swizzle parameters (a_dim0, a_dim1, etc.) are explicit because Mojo
requires them for overload resolution with parametric struct arguments.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!