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,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
__copy_ctor_is_trivial
comptime __copy_ctor_is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__move_ctor_is_trivial
comptime __move_ctor_is_trivial = True
Methods
__init__
__init__(index: IndexList[3]) -> Self
Constructor of a gemm shape record from a index tuple.
Args:
- index (
IndexList): 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](c: NDBuffer[c.dtype, 2, c.origin, c.shape, c.strides, alignment2=c.alignment2, address_space=c.address_space, exclusive=c.exclusive], a: NDBuffer[a.dtype, 2, a.origin, a.shape, a.strides, alignment2=a.alignment2, address_space=a.address_space, exclusive=a.exclusive], b: NDBuffer[b.dtype, 2, b.origin, b.shape, b.strides, alignment2=b.alignment2, address_space=b.address_space, exclusive=b.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:
static get[transpose_b: Bool, layout_c: Layout, layout_a: Layout, layout_b: Layout](c: LayoutTensor[c.dtype, layout_c, c.origin, 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[a.dtype, layout_a, a.origin, 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[b.dtype, layout_b, b.origin, 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): LayoutTensor with allocated output space. - a (
LayoutTensor): LayoutTensor containing matrix operand A. - b (
LayoutTensor): 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:
IndexList: The constructed index tuple.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!