Skip to main content
Log in

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 (Index[2]):

Implemented traits

AnyType, Copyable, ExplicitlyCopyable, Movable, UnknownDestructibility

Methods

get

static get(b: NDBuffer[b_type, 2, origin, shape], tile_n_k: Index[2]) -> Self

get_tile

get_tile[inner_size: Int](self, global_offset: GemmShape, tile_dim_nk: Index[2], valid_data_dim_nk: Index[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 (Index[2]): Tile shape based on cache size and matrix dimensions.
  • valid_data_dim_nk (Index[2]): The upper bounds for N and K dimensions.

Returns:

A view of the packed tile.

Was this page helpful?