Mojo struct
BTileGenerator
struct BTileGenerator[config: KernelConfig[config.packed_shape], a_type: DType, b_type: DType, c_type: DType, shape: DimList[shape.values], transpose_b: Bool, b_packed: Bool, origin: ImmutOrigin]
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, origin, shape, DimList.create_unknown[2]()]): - b_tile_stack_ptr (
UnsafePointer[Scalar[b_type], MutAnyOrigin]): - tile_n_k (
IndexList[2]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
Methods
get
static get(b: NDBuffer[b_type, origin, shape, DimList.create_unknown[2]()], 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, ImmutAnyOrigin, config.packed_shape, DimList.create_unknown[3]()]
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): Tile shape based on cache size and matrix dimensions. - valid_data_dim_nk (
IndexList): The upper bounds for N and K dimensions.
Returns:
NDBuffer: 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!