Skip to main content

Mojo function

output_reg_to_smem

output_reg_to_smem[BM: Int, BN: Int, WM: Int, padded_depth: Int, kv_type: DType, output_type: DType, accum_type: DType, reg_layout: Layout, o_frag_size: Int, num_consumer_threads: Int, simd_size: Int, swizzle: Swizzle, num_m_mmas: Int, num_consumer: Int, mma_thread_layout: Layout](tid: UInt32, local_warp_group_idx: UInt32, warp_x: UInt32, warp_y: UInt32, q_smem: UnsafePointer[Scalar[kv_type], address_space=AddressSpace(3), alignment=128], output_reg_tile: LayoutTensor[accum_type, reg_layout, MutableAnyOrigin, address_space=AddressSpace(5)]) -> LayoutTensor[output_type, Layout.row_major(BM, padded_depth), MutableAnyOrigin, address_space=AddressSpace(3)]

Returns:

LayoutTensor

Was this page helpful?