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,
Movable,
UnknownDestructibility
Methods
__init__
@implicit
__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__
__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[dtype, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive], a: NDBuffer[dtype, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive], b: NDBuffer[dtype, 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[dtype, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): NDBuffer with allocated output space. - a (
NDBuffer[dtype, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): NDBuffer containing matrix operand A. - b (
NDBuffer[dtype, 2, origin, shape, strides, alignment=alignment, address_space=address_space, exclusive=exclusive]): NDBuffer containing matrix operand B.
static get[transpose_b: Bool, layout_c: Layout, layout_a: Layout, layout_b: Layout](c: LayoutTensor[dtype, layout_c, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], a: LayoutTensor[dtype, layout_a, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], b: LayoutTensor[dtype, layout_b, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=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[dtype, layout_c, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): LayoutTensor with allocated output space. - a (
LayoutTensor[dtype, layout_a, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): LayoutTensor containing matrix operand A. - b (
LayoutTensor[dtype, layout_b, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): LayoutTensor 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:
The constructed index tuple.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!