Mojo struct
MMAOperandOffsetFn
struct MMAOperandOffsetFn[dtype: DType, BMN: Int, BK: Int, swizzle: TensorMapSwizzle, is_k_major: Bool, WMMA_MN: Int, WMMA_K: Int]
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
canonical_K
comptime canonical_K = (swizzle.bytes() // size_of[dtype]()) if (swizzle != TensorMapSwizzle.SWIZZLE_NONE) else BK
canonical_layout
comptime canonical_layout = tile_to_descriptor[dtype, MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].canonical_layout_flat, is_k_major]()
canonical_layout_flat
comptime canonical_layout_flat = Layout(Coord(VariadicPack(Coord(VariadicPack(Idx[8](), Idx[(BMN // 8)]())), Coord(VariadicPack(Idx[(swizzle.bytes() // size_of[dtype]())](), Idx[(MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].canonical_K // (swizzle.bytes() // size_of[dtype]()))]())))), Coord(VariadicPack(Coord(VariadicPack(Idx[(swizzle.bytes() // size_of[dtype]())](), Idx[(8 * (swizzle.bytes() // size_of[dtype]()))]())), Coord(VariadicPack(Idx[1](), Idx[0 if (MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].canonical_K == (swizzle.bytes() // size_of[dtype]())) else (BMN * (swizzle.bytes() // size_of[dtype]()))]()))))).to_layout() if is_k_major else MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].layout
canonical_layout_size
comptime canonical_layout_size = MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].canonical_layout.size()
layout
comptime layout = Layout(Coord(VariadicPack(Coord(VariadicPack(Idx[8](), Idx[(BMN // 8)]())), Coord(VariadicPack(Idx[(swizzle.bytes() // size_of[dtype]())](), Idx[(BK // (swizzle.bytes() // size_of[dtype]()))]())))), Coord(VariadicPack(Coord(VariadicPack(Idx[(swizzle.bytes() // size_of[dtype]())](), Idx[(8 * (swizzle.bytes() // size_of[dtype]()))]())), Coord(VariadicPack(Idx[1](), Idx[0 if (BK == (swizzle.bytes() // size_of[dtype]())) else (BMN * (swizzle.bytes() // size_of[dtype]()))]()))))).to_layout() if is_k_major else Layout(Coord(VariadicPack(Coord(VariadicPack(Idx[8](), Idx[(BK // 8)]())), Coord(VariadicPack(Idx[(swizzle.bytes() // size_of[dtype]())](), Idx[(BMN // (swizzle.bytes() // size_of[dtype]()))]())))), Coord(VariadicPack(Coord(VariadicPack(Idx[(swizzle.bytes() // size_of[dtype]())](), Idx[(8 * (swizzle.bytes() // size_of[dtype]()))]())), Coord(VariadicPack(Idx[1](), Idx[0 if (BMN == (swizzle.bytes() // size_of[dtype]())) else (BK * (swizzle.bytes() // size_of[dtype]()))]()))))).transpose().to_layout()
layout_size
comptime layout_size = MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].layout.size()
Methods
__init__
__init__() -> Self
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!