Skip to main content

Mojo function

shared_memory_epilogue

shared_memory_epilogue[MMA_M: Int, data_paths: Int, num_stages: Int, stage: Int, stageN: Int, c_type: DType, shared_n: Int, simd_size: Int, c_smem_upper_layout: Layout, c_smem_lower_layout: Layout, swizzle: Swizzle, compute_lambda_fn: def[dtype: DType, width: Int, *, alignment: Int = 1](IndexList[2], SIMD[dtype, width]) capturing -> SIMD[dtype, width], num_output_warps: Int](M: UInt32, N: UInt32, c_col: Int, c_row: Int, c_smem_warp_tile_upper: LayoutTensor[c_type, c_smem_upper_layout, MutAnyOrigin, address_space=c_smem_warp_tile_upper.address_space, element_layout=c_smem_warp_tile_upper.element_layout, layout_int_type=c_smem_warp_tile_upper.layout_int_type, linear_idx_type=c_smem_warp_tile_upper.linear_idx_type, masked=c_smem_warp_tile_upper.masked, alignment=c_smem_warp_tile_upper.alignment], c_smem_warp_tile_lower: LayoutTensor[c_type, c_smem_lower_layout, MutAnyOrigin, address_space=c_smem_warp_tile_lower.address_space, element_layout=c_smem_warp_tile_lower.element_layout, layout_int_type=c_smem_warp_tile_lower.layout_int_type, linear_idx_type=c_smem_warp_tile_lower.linear_idx_type, masked=c_smem_warp_tile_lower.masked, alignment=c_smem_warp_tile_lower.alignment])

shared_memory_epilogue[MMA_M: Int, data_paths: Int, num_stages: Int, stage: Int, stageN: Int, c_type: DType, shared_n: Int, simd_size: Int, swizzle: Swizzle, compute_lambda_fn: def[dtype: DType, width: Int, *, alignment: Int = 1](IndexList[2], SIMD[dtype, width]) capturing -> SIMD[dtype, width], num_output_warps: Int](M: UInt32, N: UInt32, c_col: Int, c_row: Int, c_smem_warp_tile_upper: TileTensor[c_type, c_smem_warp_tile_upper.LayoutType, c_smem_warp_tile_upper.origin, address_space=c_smem_warp_tile_upper.address_space, linear_idx_type=c_smem_warp_tile_upper.linear_idx_type, element_size=c_smem_warp_tile_upper.element_size], c_smem_warp_tile_lower: TileTensor[c_type, c_smem_warp_tile_lower.LayoutType, c_smem_warp_tile_lower.origin, address_space=c_smem_warp_tile_lower.address_space, linear_idx_type=c_smem_warp_tile_lower.linear_idx_type, element_size=c_smem_warp_tile_lower.element_size])

Was this page helpful?