Skip to main content

Mojo function

pack_b

pack_b[transpose_b: Bool, simd_size: Int, inner_size: Int, a_type: DType, b_type: DType, c_type: DType, src_shape: DimList[src_shape.values], dst_shape: DimList[dst_shape.values]](dst: NDBuffer[b_type, dst.origin, dst_shape, DimList.create_unknown[2]()], src: NDBuffer[b_type, src.origin, src_shape, DimList.create_unknown[2]()], tile_n: Int, tile_k: Int)

Utility function to pack the entire B matrix, such that each [tile_n // inner_size, tile_k, inner_size] tile of src is contiguous in dst.

Tiles (not tile contents) are stored in row major order, so tile[i, j] is tile_n * tile_k bytes away from tile[i, j+1].

Was this page helpful?