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 function

fma

def fma[mode: StringSlice[ImmStaticOrigin], dtype: DType](z_row_index: Int, x_row_index: Int, y_row_index: Int, clear_z: Bool)

Issues an AMX fused multiply-add of X and Y register rows accumulated into the Z accumulator.

Supports two modes selected by the mode parameter. In "ROW" mode the instruction performs an elementwise fused multiply-add computing Z[z_row_index][:] += X[x_row_index][:] * Y[y_row_index][:]. In "TILE" mode it computes the outer product Y[y_row_index][:] X X[x_row_index][:] (a 16x16 matrix) and accumulates the result into Z[z_row_index::step 4][:]. When clear_z is true the existing Z values are ignored instead of accumulated. Issues the fma.fp32 AMX instruction.

Parameters:

mode (StringSlice[ImmStaticOrigin]): Either "ROW" for elementwise fma or "TILE" for outer-product fma.
dtype (DType): Must be DType.float32.

Args:

z_row_index (Int): Row of Z to write; must be in [0, 8) in row mode or [0, 4) in tile mode.
x_row_index (Int): Row of X to read; must be in [0, 8).
y_row_index (Int): Row of Y to read; must be in [0, 8).
clear_z (Bool): When true, clears the target Z rows instead of accumulating into them.