IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

AppleM5MatMul

struct AppleM5MatMul[in_type: DType, c_type: DType = DType.float32, transpose_b: Bool = False, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, block_m: Int = Int(64), block_n: Int = Int(64), block_k: Int = Int(16), sg_m: Int = Int(32), sg_n: Int = Int(32), k_unroll: Int = Int(1), use_x2: Bool = False, linear_idx_type: DType = DType.int64, clamp_edge: Bool = False, b_type: DType = in_type, accum_type: DType = DType.float32]

Apple M5 simdgroup-tiled GEMM (Metal 4 hardware MMA).

run is the GPU kernel entry (TileTensor operands; M/N/K derived from them). run_split_k_partial / run_split_k_reduce are the split-K kernels (bf16 family; split-K is not wired for b_type != in_type). Launch via enqueue_apple_matmul / enqueue_apple_matmul_split_k; the FP8 W8A16 dense launcher is enqueue_matmul2d_fp8.

Parameters​

  • ​in_type (DType): A element type (fp16, bf16, fp32; int8 at the MMA-op level).
  • ​c_type (DType): Output element type (fp16, bf16, fp32). Accumulation is accum_type.
  • ​transpose_b (Bool): If True, B is (N, K) row-major (viewed as col_major(K, N)); otherwise B is (K, N) row-major.
  • ​elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional fused epilogue; receives SIMD[c_type, width] at absolute (row, col) (AMD's contract).
  • ​block_m (Int): Threadgroup block rows BM (M tile). Multiple of SG_M.
  • ​block_n (Int): Threadgroup block cols BN (N tile). Multiple of SG_N.
  • ​block_k (Int): K-strip depth BK per accumulate step. Multiple of MMA_K (16). The conv path (enqueue_apple_conv2d) overrides these to tune the fused im2col GEMM independently of the dense GEMM defaults.
  • ​sg_m (Int): Simdgroup subtile rows SG_M. Multiple of MMA_M (16); fixes the M MMA-fragment count NUM_MMA_M = SG_M / 16.
  • ​sg_n (Int): Simdgroup subtile cols SG_N. Multiple of MMA_N (16); fixes the N MMA-fragment count NUM_MMA_N = SG_N / 16.
  • ​k_unroll (Int): Dense run path interior K-loop unroll factor (default 1, no unroll). Processes k_unroll BK-strips per loop pass, amortizing loop-control overhead across the group. Regresses severely if combined with use_x2 (register pressure) -- the dispatcher never sets both.
  • ​use_x2 (Bool): NT bf16/fp16 "double-strip" dense MMA (block_k=32, MmaOpApple.mma_dense_x2): both operands are K-contiguous, so one call covers 32 K, halving the K-loop iteration count.
  • ​linear_idx_type (DType): Integer type for the dense run path's A/B/C TileTensor offset arithmetic (Apple's scalar ALU is faster on 32-bit). Only safe when every offset provably fits int32.
  • ​clamp_edge (Bool): Enables clamp_v2: shift a ragged last M/N tile to m-BM/n-BN for a full, fast load, then store back only its owned region so it doesn't overwrite the previous tile. Dense fp32 only, no elementwise_lambda_fn. Only run_chained consumes it; run/run_conv never set it.
  • ​b_type (DType): B/weight operand dtype fed to the MMA. Defaults to in_type (the dense/conv case). The W8A16 FP8 path sets b_type=float8_e4m3fn: the A slab stays in_type (bf16) and only the B fragment widens to the native-fp8 MMA operand (the weight-loader seam swap). See WeightLoader / enqueue_matmul2d_fp8.
  • ​accum_type (DType): MMA accumulator dtype. fp32 for dense/fp8; a later int8 slice uses int32.

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

BK​

comptime BK = block_k

BM​

comptime BM = block_m

BN​

comptime BN = block_n

Mma​

comptime Mma = MmaOpApple[accum_type, in_type, AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].NUM_MMA_M, AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].NUM_MMA_N, b_type=b_type]

NUM_MMA_M​

comptime NUM_MMA_M = (AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].SG_M // Int(16))

NUM_MMA_N​

comptime NUM_MMA_N = (AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].SG_N // Int(16))

NUM_SG​

comptime NUM_SG = (AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].NUM_SG_M * AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].NUM_SG_N)

NUM_SG_M​

comptime NUM_SG_M = (AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].BM // AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].SG_M)

NUM_SG_N​

comptime NUM_SG_N = (AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].BN // AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].SG_N)

REDUCE_BLOCK​

comptime REDUCE_BLOCK = 256

SG_M​

comptime SG_M = sg_m

SG_N​

comptime SG_N = sg_n

THREADS_PER_BLOCK​

comptime THREADS_PER_BLOCK = (AppleM5MatMul[in_type, c_type, transpose_b, elementwise_lambda_fn, block_m, block_n, block_k, sg_m, sg_n, k_unroll, use_x2, linear_idx_type, clamp_edge, b_type, accum_type].NUM_SG * _resolve_warp_size())

Methods​

morton_decode_2d​

static def morton_decode_2d(flat_idx: UInt32) -> Tuple[UInt32, UInt32]

Decode a linear index to (tile_m, tile_n) via Morton Z-order.

Even bits of flat_idx -> tile_n, odd bits -> tile_m. The decoded pair may fall outside any rectangular grid that isn't a power-of-2 square; the caller checks bounds.

Future: as of AIR 4.1 a bit_interleave intrinsic can do this interleave directly, replacing the hand-rolled bit-scatter below. A Gilbert-curve dispatch order was also explored for better locality, but it needs a per-shape predispatch kernel to generate the order.

Args:

  • ​flat_idx (UInt32): The linear threadgroup index to split into interleaved Z-order bits; even bits form tile_n, odd bits form tile_m.

Returns:

Tuple[UInt32, UInt32]

morton_decode_2d_rect​

static def morton_decode_2d_rect(flat_idx: UInt32, log2_m: UInt32, log2_n: UInt32) -> Tuple[UInt32, UInt32]

Decode flat_idx to (tile_m, tile_n) over a (1<<log2_m) x (1<<log2_n) grid.

Z-order covers a min(side_m, side_n) square core; remaining bits sweep the longer axis. Reduces to morton_decode_2d when log2_m == log2_n.

Args:

  • ​flat_idx (UInt32): The linear threadgroup index to decode into a 2-D tile coordinate.
  • ​log2_m (UInt32): Base-2 logarithm of the grid's M extent; the decoded tile_m ranges over [0, 1<<log2_m).
  • ​log2_n (UInt32): Base-2 logarithm of the grid's N extent; the decoded tile_n ranges over [0, 1<<log2_n).

Returns:

Tuple[UInt32, UInt32]

run​

static def run[c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, c_storage: TensorStorage, a_storage: TensorStorage, b_storage: TensorStorage](c: TileTensor[c_type, c_layout, MutAnyOrigin, Storage=c_storage], a: TileTensor[in_type, a_layout, ImmutAnyOrigin, Storage=a_storage], b: TileTensor[b_type, b_layout, ImmutAnyOrigin, Storage=b_storage], log2_grid_m: UInt32, log2_grid_n: UInt32)

GEMM kernel entry; M/N/K derive from the operands.

A is (M, K) row-major and C is (M, N) row-major; B is (K, N) for transpose_b=False or (N, K) for transpose_b=True. Grid is (1<<log2_grid_m) * (1<<log2_grid_n) threadgroups of 128 threads; OOB threadgroups early-return after Morton decode.

Thin wrapper over _run_gemm_body: derives K, pre-tiles this simdgroup's (SG_M, K) A slab, and constructs the DenseALoader that reads it.

Parameters:

  • ​c_layout (TensorLayout): TensorLayout of the output C operand.
  • ​a_layout (TensorLayout): TensorLayout of the A operand.
  • ​b_layout (TensorLayout): TensorLayout of the B operand.
  • ​c_storage (TensorStorage): TensorStorage of the output C operand.
  • ​a_storage (TensorStorage): TensorStorage of the A operand.
  • ​b_storage (TensorStorage): TensorStorage of the B operand.

Args:

run_chained​

static def run_chained[c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, c_storage: TensorStorage, a_storage: TensorStorage, b_storage: TensorStorage, seed_from_output: Bool = False](c: TileTensor[c_type, c_layout, MutAnyOrigin, Storage=c_storage], a: TileTensor[in_type, a_layout, ImmutAnyOrigin, Storage=a_storage], b: TileTensor[b_type, b_layout, ImmutAnyOrigin, Storage=b_storage], log2_grid_m: UInt32, log2_grid_n: UInt32, k_strip_start: Int, k_strip_end: Int)

clamp_v2 chained-pass kernel. enqueue_apple_matmul_clamp_chain launches it twice: pass 0 zero-seeds [k_strip_start, k_strip_end), pass 1 seeds from c and covers the rest, so no partials buffer or reduce kernel is needed. Otherwise identical to run; requires Self.clamp_edge=True (asserted below).

run_conv​

static def run_conv[c_layout: TensorLayout, input_layout: TensorLayout, b_layout: TensorLayout, c_storage: TensorStorage, input_storage: TensorStorage, b_storage: TensorStorage, c_aligned: Bool = False](c: TileTensor[c_type, c_layout, MutAnyOrigin, Storage=c_storage], input: TileTensor[in_type, input_layout, ImmutAnyOrigin, Storage=input_storage], b: TileTensor[in_type, b_layout, ImmutAnyOrigin, Storage=b_storage], conv: ConvIm2colParams, log2_grid_m: UInt32, log2_grid_n: UInt32)

Fused conv2d GEMM entry; M/N/K derive from C and the conv params.

C is (M, N) row-major with M = N_batch*H_out*W_out, N = C_out. input is the 4-D NHWC source (flat pointer used for the gather). B is the filter (N, K) with transpose_b=True (the NK layout the existing matmul path uses). Requires transpose_b == True. Grid is identical to run. The fused epilogue (if any) and the partial-edge bounds are the same as run -- only the A-fragment producer changes.

Thin wrapper over _run_gemm_body: derives K = R*S*C and constructs an Im2colALoader over input (no slab -- the gather reads NHWC directly). conv is threaded to the body as a value arg (NOT held in the loader -- see Im2colALoader for the Metal addrspace reason), then runs the body.

Parameters:

  • ​c_layout (TensorLayout): TensorLayout of the output C operand.
  • ​input_layout (TensorLayout): TensorLayout of the NHWC input operand.
  • ​b_layout (TensorLayout): TensorLayout of the filter B operand.
  • ​c_storage (TensorStorage): TensorStorage of the output C operand.
  • ​input_storage (TensorStorage): TensorStorage of the NHWC input operand.
  • ​b_storage (TensorStorage): TensorStorage of the filter B operand.
  • ​c_aligned (Bool): When True, assume conv.C is a multiple of 8 so the channel run is contiguous and a single width-8 load suffices (defaults to False).

Args:

run_split_k_partial​

static def run_split_k_partial[a_layout: TensorLayout, b_layout: TensorLayout, a_storage: TensorStorage, b_storage: TensorStorage](partials_ptr: Pointer[Float32, MutAnyOrigin, _safe=False], a: TileTensor[in_type, a_layout, ImmutAnyOrigin, Storage=a_storage], b: TileTensor[in_type, b_layout, ImmutAnyOrigin, Storage=b_storage], log2_grid_m: UInt32, log2_grid_n: UInt32, k_per_split: Int)

One 64x64 output tile's fp32 partial over a BK-aligned K-slice.

Grid is (side_m * side_n) * num_splits threadgroups. The split index selects the K-range [s*k_per_split, min(K, (s+1)*k_per_split)) and the partials[s] output matrix; the tile index is Morton-decoded as in run. k_per_split is a multiple of BK, so every split but the last is full BK strips; the last may carry a partial-BK tail. No cast, no epilogue -- raw fp32 accumulator out.

Parameters:

Args:

run_split_k_reduce​

static def run_split_k_reduce[c_layout: TensorLayout, c_storage: TensorStorage](c: TileTensor[c_type, c_layout, MutAnyOrigin, Storage=c_storage], partials_ptr: Pointer[Float32, ImmutAnyOrigin, _safe=False], num_splits: Int)

Sum num_splits fp32 partials per output element, cast, store / fuse.

One thread per output element; idx = block_idx.x * block_dim.x + thread_idx.x. The fused elementwise_lambda_fn (if any) sees the absolute (row, col) and the final SIMD[c_type, 1].

Parameters:

  • ​c_layout (TensorLayout): TensorLayout of the output C operand.
  • ​c_storage (TensorStorage): TensorStorage of the output C operand.

Args:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact