Skip to main content

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?