Mojo struct
GemmShape
struct GemmShape
Helper class to unpack gemm dimension and layout.
Fields
- M (
Int): - N (
Int): - K (
Int):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
__init__(index: IndexList[3]) -> Self
Constructor of a gemm shape record from a index tuple.
Args:
- index (
IndexList[3]): The int tuple containing the index(m,n,k).
__getitem__
__setitem__
__setitem__(mut self, idx: Int, value: Int)
__add__
__add__(self, rhs: Self) -> Self
Coordinate-wise addition of two gemm shape records.
Args:
- rhs (
Self): Another gemm shape record to add with.
__sub__
__sub__(self, rhs: Self) -> Self
Coordinate-wise subtraction of two gemm shape records.
Args:
- rhs (
Self): Another gemm shape record to subtract with.
get
static get[transpose_b: Bool, layout_c: Layout, layout_a: Layout, layout_b: Layout](c: LayoutTensor[layout_c, address_space=c.address_space, element_layout=c.element_layout, layout_int_type=c.layout_int_type, linear_idx_type=c.linear_idx_type, masked=c.masked, alignment=c.alignment], a: LayoutTensor[layout_a, address_space=a.address_space, element_layout=a.element_layout, layout_int_type=a.layout_int_type, linear_idx_type=a.linear_idx_type, masked=a.masked, alignment=a.alignment], b: LayoutTensor[layout_b, address_space=b.address_space, element_layout=b.element_layout, layout_int_type=b.layout_int_type, linear_idx_type=b.linear_idx_type, masked=b.masked, alignment=b.alignment]) -> Self
Constructor of a gemm shape record from input buffers.
M, N, and K are intentionally calculated using a and c ONLY. This
is because b may be padded to a multiple of the tile size if it has
been pre-packed.
Args:
- c (
LayoutTensor[layout_c, address_space=c.address_space, element_layout=c.element_layout, layout_int_type=c.layout_int_type, linear_idx_type=c.linear_idx_type, masked=c.masked, alignment=c.alignment]): LayoutTensor with allocated output space. - a (
LayoutTensor[layout_a, address_space=a.address_space, element_layout=a.element_layout, layout_int_type=a.layout_int_type, linear_idx_type=a.linear_idx_type, masked=a.masked, alignment=a.alignment]): LayoutTensor containing matrix operand A. - b (
LayoutTensor[layout_b, address_space=b.address_space, element_layout=b.element_layout, layout_int_type=b.layout_int_type, linear_idx_type=b.linear_idx_type, masked=b.masked, alignment=b.alignment]): LayoutTensor containing matrix operand B.
static get[transpose_b: Bool](c: TileTensor[address_space=c.address_space, linear_idx_type=c.linear_idx_type, element_size=c.element_size], a: TileTensor[address_space=a.address_space, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[address_space=b.address_space, linear_idx_type=b.linear_idx_type, element_size=b.element_size]) -> Self
Constructor of a gemm shape record from TileTensor inputs.
M, N, and K are intentionally calculated using a and c ONLY. This
is because b may be padded to a multiple of the tile size if it has
been pre-packed.
Args:
- c (
TileTensor[address_space=c.address_space, linear_idx_type=c.linear_idx_type, element_size=c.element_size]): TileTensor with allocated output space. - a (
TileTensor[address_space=a.address_space, linear_idx_type=a.linear_idx_type, element_size=a.element_size]): TileTensor containing matrix operand A. - b (
TileTensor[address_space=b.address_space, linear_idx_type=b.linear_idx_type, element_size=b.element_size]): TileTensor containing matrix operand B.
as_index
as_index(self) -> IndexList[3]
Utility to convert the underlying data to an index tuple. So that the utilities such as elementwise add can be used.
Returns:
IndexList[3]: The constructed index tuple.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!