Mojo struct
BlockScaledTileStorage
struct BlockScaledTileStorage[a_type: DType, b_type: DType, c_type: DType, sfa_type: DType, sfb_type: DType, a_dim0: Int, a_dim1: Int, b_dim0: Int, b_dim1: Int, c_dim0: Int, c_dim1: Int, sfa_dim0: Int, sfa_dim1: Int, sfb_dim0: Int, sfb_dim1: Int, num_pipeline_stages: Int, num_output_stages: Int]
Storage for block-scaled matmul tiles (A, B, C, SFA, SFB).
Single source of truth for block-scaled tile arrays and storage.
All tiles use TileTensor natively. C tiles also store as TileTensor but provide a LayoutTensor accessor for epilogue_components.mojo compatibility with .reshape[] and .tile[] methods.
IMPORTANT: Field order preserves SMEM layout compatibility: a, b, c, sfa, sfb.
Parameters
- a_type (
DType): Data type for A matrix tiles. - b_type (
DType): Data type for B matrix tiles. - c_type (
DType): Data type for C matrix tiles. - sfa_type (
DType): Data type for A scale factor tiles. - sfb_type (
DType): Data type for B scale factor tiles. - a_dim0 (
Int): First dimension for A tiles. - a_dim1 (
Int): Second dimension for A tiles. - b_dim0 (
Int): First dimension for B tiles. - b_dim1 (
Int): Second dimension for B tiles. - c_dim0 (
Int): First dimension for C tiles (OutputM). - c_dim1 (
Int): Second dimension for C tiles (OutputN). - sfa_dim0 (
Int): First dimension for SFA tiles. - sfa_dim1 (
Int): Second dimension for SFA tiles. - sfb_dim0 (
Int): First dimension for SFB tiles. - sfb_dim1 (
Int): Second dimension for SFB tiles. - num_pipeline_stages (
Int): Number of input pipeline stages. - num_output_stages (
Int): Number of output pipeline stages.
Fields
- a_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].ATileArray.Storage): - b_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].BTileArray.Storage): - c_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].CTileArray.Storage): - sfa_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].SFATileArray.Storage): - sfb_tiles_storage (
BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].SFBTileArray.Storage):
Implemented traits
AnyType,
ImplicitlyDestructible
comptime members
__del__is_trivial
comptime __del__is_trivial = True
ATileArray
comptime ATileArray = SMemTileArray2D[a_type, a_dim0, a_dim1, num_pipeline_stages]
BTileArray
comptime BTileArray = SMemTileArray2D[b_type, b_dim0, b_dim1, num_pipeline_stages]
c_tile_layout
comptime c_tile_layout = Layout.row_major(c_dim0, c_dim1)
CTileArray
comptime CTileArray = SMemTileArray2DRowMajor[c_type, c_dim0, c_dim1, num_output_stages]
CTileArrayLT
comptime CTileArrayLT = SMemTileArray[c_type, BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].c_tile_layout, num_output_stages, 128]
sfa_layout
comptime sfa_layout = Layout[Coord[ComptimeInt[32], ComptimeInt[(sfa_dim0 // 32)]], Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfa_dim1 // 16)]], Coord[ComptimeInt[16], ComptimeInt[(sfa_dim1 * 32)]], Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]]](Coord[Coord[ComptimeInt[32], ComptimeInt[(sfa_dim0 // 32)]], Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfa_dim1 // 16)]]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[32], ComptimeInt[(sfa_dim0 // 32)]], Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfa_dim1 // 16)]]](Coord[ComptimeInt[32], ComptimeInt[(sfa_dim0 // 32)]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[32], ComptimeInt[(sfa_dim0 // 32)]](Idx[32](), Idx[(sfa_dim0 // 32)]())), Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfa_dim1 // 16)]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfa_dim1 // 16)]](Coord[ComptimeInt[4], ComptimeInt[4]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[4], ComptimeInt[4]](Idx[4](), Idx[4]())), Idx[(sfa_dim1 // 16)]())))), Coord[Coord[ComptimeInt[16], ComptimeInt[(sfa_dim1 * 32)]], Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[16], ComptimeInt[(sfa_dim1 * 32)]], Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]]](Coord[ComptimeInt[16], ComptimeInt[(sfa_dim1 * 32)]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[16], ComptimeInt[(sfa_dim1 * 32)]](Idx[16](), Idx[(sfa_dim1 * 32)]())), Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]](Coord[ComptimeInt[1], ComptimeInt[4]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[1], ComptimeInt[4]](Idx[1](), Idx[4]())), Idx[512]())))))
SFATileArray
comptime SFATileArray = SMemTileArrayWithLayout[sfa_type, BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].sfa_layout, num_pipeline_stages]
sfb_layout
comptime sfb_layout = Layout[Coord[ComptimeInt[32], ComptimeInt[(sfb_dim0 // 32)]], Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfb_dim1 // 16)]], Coord[ComptimeInt[16], ComptimeInt[(sfb_dim1 * 32)]], Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]]](Coord[Coord[ComptimeInt[32], ComptimeInt[(sfb_dim0 // 32)]], Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfb_dim1 // 16)]]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[32], ComptimeInt[(sfb_dim0 // 32)]], Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfb_dim1 // 16)]]](Coord[ComptimeInt[32], ComptimeInt[(sfb_dim0 // 32)]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[32], ComptimeInt[(sfb_dim0 // 32)]](Idx[32](), Idx[(sfb_dim0 // 32)]())), Coord[Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfb_dim1 // 16)]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[4], ComptimeInt[4]], ComptimeInt[(sfb_dim1 // 16)]](Coord[ComptimeInt[4], ComptimeInt[4]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[4], ComptimeInt[4]](Idx[4](), Idx[4]())), Idx[(sfb_dim1 // 16)]())))), Coord[Coord[ComptimeInt[16], ComptimeInt[(sfb_dim1 * 32)]], Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[16], ComptimeInt[(sfb_dim1 * 32)]], Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]]](Coord[ComptimeInt[16], ComptimeInt[(sfb_dim1 * 32)]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[16], ComptimeInt[(sfb_dim1 * 32)]](Idx[16](), Idx[(sfb_dim1 * 32)]())), Coord[Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]](VariadicPack[True, MutExternalOrigin, True, CoordLike, Coord[ComptimeInt[1], ComptimeInt[4]], ComptimeInt[512]](Coord[ComptimeInt[1], ComptimeInt[4]](VariadicPack[True, MutExternalOrigin, True, CoordLike, ComptimeInt[1], ComptimeInt[4]](Idx[1](), Idx[4]())), Idx[512]())))))
SFBTileArray
comptime SFBTileArray = SMemTileArrayWithLayout[sfb_type, BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].sfb_layout, num_pipeline_stages]
Methods
a_tiles
a_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].ATileArray
Get A tile array accessor.
Returns:
BlockScaledTileStorage
b_tiles
b_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].BTileArray
Get B tile array accessor.
Returns:
BlockScaledTileStorage
c_tiles
c_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].CTileArrayLT
Get C tile array accessor (LayoutTensor-based for backward compat).
Returns LayoutTensor view for compatibility with tile_writer.mojo which uses .reshape[] and .tile[] methods.
Returns:
BlockScaledTileStorage
c_tiles_tt
c_tiles_tt(ref[AddressSpace._value._mlir_value] self) -> BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].CTileArray
Get C tile array accessor (TileTensor-based).
Returns native TileTensor for future TileTensor-native code paths.
Returns:
BlockScaledTileStorage
sfa_tiles
sfa_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].SFATileArray
Get SFA tile array accessor.
Returns:
BlockScaledTileStorage
sfb_tiles
sfb_tiles(ref[AddressSpace._value._mlir_value] self) -> BlockScaledTileStorage[a_type, b_type, c_type, sfa_type, sfb_type, a_dim0, a_dim1, b_dim0, b_dim1, c_dim0, c_dim1, sfa_dim0, sfa_dim1, sfb_dim0, sfb_dim1, num_pipeline_stages, num_output_stages].SFBTileArray
Get SFB tile array accessor.
Returns:
BlockScaledTileStorage
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!