For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

BTileGenerator

struct BTileGenerator[config: KernelConfig, a_type: DType, b_type: DType, c_type: DType, b_layout: TensorLayout, transpose_b: Bool, b_packed: Bool, origin: ImmOrigin]

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.

Parameters

• ​config (KernelConfig): Kernel configuration supplying simd_size and kernel_cols used by the packing routines.
• ​a_type (DType): Element type of the A operand of the matmul.
• ​b_type (DType): Element type of the B operand being packed.
• ​c_type (DType): Element type of the C output of the matmul.
• ​b_layout (TensorLayout): Layout of the B TileTensor.
• ​transpose_b (Bool): True if the B operand is transposed, stored as [N, K].
• ​b_packed (Bool): True if B is already pre-packed into the mlas layout.
• ​origin (ImmOrigin): Origin of the B TileTensor.

Fields

• ​b (TileTensor[b_type, b_layout, origin]):
• ​b_tile_stack_ptr (Pointer[Scalar[b_type], MutUntrackedOrigin]):
• ​tile_n_k (IndexList[Int(2)]):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable

comptime members

PackedTileLayout

comptime PackedTileLayout = Layout[*?, *?]

Methods

get

static def get(b: TileTensor[b_type, b_layout, origin], tile_n_k: IndexList[Int(2)]) -> Self

get_tile

def get_tile[inner_size: Int](self, global_offset: GemmShape, tile_dim_nk: IndexList[Int(2)], valid_data_dim_nk: IndexList[Int(2)]) -> TileTensor[b_type, Layout[*?, *?], ImmutAnyOrigin]

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.

Parameters:

• ​inner_size (Int): Size of the inner dimension along N in the packed tile; must be a multiple of simd_size.

Args:

• ​global_offset (GemmShape): Offset in the global M, N, K dimensions.
• ​tile_dim_nk (IndexList[Int(2)]): Tile shape based on cache size and matrix dimensions.
• ​valid_data_dim_nk (IndexList[Int(2)]): The upper bounds for N and K dimensions.

Returns:

TileTensor[b_type, Layout[*?, *?], ImmutAnyOrigin]: A view of the packed tile.