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: 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 (
TileTensor[b_type, b_layout, origin]): - b_tile_stack_ptr (
UnsafePointer[Scalar[b_type], MutAnyOrigin]): - tile_n_k (
IndexList[2]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
comptime members
PackedTileLayout
comptime PackedTileLayout = Layout[*?, *?]
Methods
get
static get(b: TileTensor[b_type, b_layout, origin], 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]) -> 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.
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:
TileTensor[b_type, Layout[*?, *?], ImmutAnyOrigin]: 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!