Mojo struct
GemmShape
@register_passable(trivial)
struct GemmShape
Helper class to unpack gemm dimension and layout.
Fields
- M (
Int): - N (
Int): - K (
Int):
Implemented traits
AnyType,
Copyable,
ExplicitlyCopyable,
Movable,
UnknownDestructibility
Methods
__init__
@implicit
__init__(index: Index[3]) -> Self
Constructor of a gemm shape record from a index tuple.
Args:
- index (
Index[3]): The int tuple containing the index(m,n,k).
__getitem__
__getitem__(self, idx: Int) -> Int
__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](c: NDBuffer[type, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive], a: NDBuffer[type, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive], b: NDBuffer[type, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]) -> 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 (
NDBuffer[type, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): NDBuffer with allocated output space. - a (
NDBuffer[type, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): NDBuffer containing matrix operand A. - b (
NDBuffer[type, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): NDBuffer containing matrix operand B.
as_index
as_index(self) -> Index[3]
Utility to convert the underlying data to an index tuple. So that the utilities such as elementwise add can be used.
Returns:
The constructed index tuple.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!