For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
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,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
def __init__(index: IndexList[Int(3)]) -> Self
Constructor of a gemm shape record from a index tuple.
Args:
- index (
IndexList[Int(3)]): The int tuple containing the index(m,n,k).
__getitem__
__setitem__
def __setitem__(mut self, idx: Int, value: Int)
__add__
def __add__(self, rhs: Self) -> Self
Coordinate-wise addition of two gemm shape records.
Args:
- rhs (
Self): Another gemm shape record to add with.
__sub__
def __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 def 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 def get[transpose_b: Bool](c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type, element_size=c.element_size], a: TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[Storage=b.Storage, 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[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type, element_size=c.element_size]): TileTensor with allocated output space. - a (
TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type, element_size=a.element_size]): TileTensor containing matrix operand A. - b (
TileTensor[Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type, element_size=b.element_size]): TileTensor containing matrix operand B.
as_index
def as_index(self) -> IndexList[Int(3)]
Utility to convert the underlying data to an index tuple. So that the utilities such as elementwise add can be used.
Returns:
IndexList[Int(3)]: The constructed index tuple.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!