Skip to main content

Mojo function

multi_stage_store_C

multi_stage_store_C[c_type: DType, c_smem_layout: Layout, c_layout: Layout, c_tensor_layout: Layout, c_desc_layout: Layout, num_accum_pipeline_stages: UInt, /, *, accum_type: DType, block_tile_shape: IndexList[3], mma_shape: IndexList[3], stage_stride_cols: UInt, c_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, cta_group: Int = 1, num_output_warps: UInt = 4, max_tmem_cols: UInt = 512, elementwise_lambda_fn: OptionalReg[elementwise_epilogue_type] = None, transpose_c: Bool = False](c_iter: LayoutTensorIter[c_type, c_smem_layout, MutAnyOrigin, address_space=AddressSpace.SHARED, alignment=128], c_tma_op: TMATensorTile[c_type, c_layout, c_desc_layout], c: LayoutTensor[c_type, c_tensor_layout, MutAnyOrigin], accum_pipeline_consumer_state: PipelineState[num_accum_pipeline_stages], accum_full_mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], accum_empty_mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], tmem_addr: UInt32, work_tile_coord: Tuple[UInt, UInt], group_end_idx: UInt32, elect_one_warp: Bool, M: UInt32, N: UInt32)

Was this page helpful?