Mojo struct
BTileGenerator
struct BTileGenerator[mut: Bool, //, config: KernelConfig, a_type: DType, b_type: DType, c_type: DType, shape: DimList, transpose_b: Bool, b_packed: Bool, origin: Origin[mut]]
Struct to encapsulate a tile of B that supports prepacking.
If b_packed is true, calls to get_tile will return a buffer view from B. Otherwise, calls to get_tile will copy a tile from B into a stack allocated scratch buffer and return a view of that.
Fields
- b (
NDBuffer[b_type, 2, origin, shape]): - b_tile_stack_ptr (
UnsafePointer[SIMD[b_type, 1]]): - tile_n_k (
IndexList[2]):
Implemented traits
AnyType,
Copyable,
ExplicitlyCopyable,
Movable,
UnknownDestructibility
Methods
get
static get(b: NDBuffer[b_type, 2, origin, shape], tile_n_k: IndexList[2]) -> Self
get_tile
get_tile[inner_size: Int](self, global_offset: GemmShape, tile_dim_nk: IndexList[2], valid_data_dim_nk: IndexList[2]) -> NDBuffer[b_type, 3, MutableAnyOrigin, config.packed_shape]
Get a packed matrix (B) tile.
valid_data_tile_nk is ignored for pre-packing, where the tile is padded to have shape of tile_dim_nk.
Args:
- global_offset (
GemmShape): Offset in the global M, N, K dimensions. - tile_dim_nk (
IndexList[2]): Tile shape based on cache size and matrix dimensions. - valid_data_dim_nk (
IndexList[2]): The upper bounds for N and K dimensions.
Returns:
A view of the packed tile.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!