Mojo struct
TiledMatmul
struct TiledMatmul[config: KernelConfig[config.packed_shape], transpose_b: Bool, b_packed: Bool, elementwise_epilogue_enabled: Bool, kernel_id: InnerKernelID, a_type: DType, a_shape: DimList[a_shape.values], a_origin: ImmutOrigin, b_type: DType, b_shape: DimList[b_shape.values], b_origin: ImmutOrigin, c_type: DType, c_shape: DimList[c_shape.values], c_origin: MutOrigin, algorithm: InnerMatmulKernel]
Tiled matmul implementation integrating packing, inner loop and tile partitions.
TODO: add tag based implementation dispatch. TODO: add fusion hooks.
Fields
- alg (
algorithm): - c (
NDBuffer[c_type, c_origin, c_shape, DimList.create_unknown[2]()]): - a (
NDBuffer[a_type, a_origin, a_shape, DimList.create_unknown[2]()]): - b (
NDBuffer[b_type, b_origin, b_shape, DimList.create_unknown[2]()]): - tile_n_k (
IndexList[2]): - global_tile_offset (
GemmShape): - global_tile_shape (
GemmShape): - b_tile_generator (
BTileGenerator[config, a_type, b_type, c_type, b_shape, transpose_b, b_packed, b_origin]): - elementwise_epilogue_fn (
def(GemmShape, GemmShape) escaping -> None):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!