DOCSEARCHbyAlgolia
Submit Docs

corsix/amx

repository·main·Indexed 22 days ago

https://github.com/corsix/amx

Low-level access, documentation, and reference implementations for Apple's undocumented AMX (Apple Matrix Extensions) instructions used on M1, M2, M3, and M4 series chips. Includes the aarch64.h header for instruction access and detailed specifications for matrix and vector computations, including the extrh and extrv instructions, hardware variations across Apple Silicon generations, and emulation logic for EXTRX.

Tokens
24.3K
Snippets
23
Records
90
Agent score
77%

What's inside corsix-amx

  1. Overview of Apple AMX (Apple Matrix Extensions)

    main

    Apple AMX is a set of undocumented and unsupported instructions used on M1, M2, M3, and M4 series chips to perform heavy matrix and vector computations via a special accelerator execution unit. It is distinct from Intel's AMX.

    Hardware Model

    AMX operates on a grid of compute units (e.g., a 32x32 grid) that can perform multiply-accumulate operations. It uses two pools of registers:

    • X registers: Contain 32 16-bit elements, 16 32-bit elements, or 8 64-bit elements.
    • Y registers: Contain 32 16-bit elements, 16 32-bit elements, or 8 64-bit elements.

    A single instruction can perform a full outer product by multiplying every element of an X register with every element of a Y register and accumulating the result into a Z element.

    Supported Data Types

    • IEEE754: f16, f32, or f64 (operands must have the same width).
    • Mixed Precision: f16 multiplicands accumulating onto f32.
    • M2 Specific: bf16 multiplicands accumulating onto bf16 or f32.
    • Integer: 8-bit or 16-bit multiplicands accumulating onto 16-bit or 32-bit integers (various signednesses).

  2. AMX Load and Store Instructions Summary

    main

    AMX (Apple Matrix Extension) provides several instructions for moving data between memory and AMX registers (X, Y, and Z). These instructions support single register operations, pairs, and multiple register loads/stores depending on the hardware generation (M1, M2, or M3).

    Instruction Set Overview

    InstructionThemeDescription
    ldxLoad Xx[i] = memory[i]
    ldyLoad Yy[i] = memory[i]
    ldz / ldziLoad Zz[_][i] = memory[i] (Interleaved Z)
    stxStore Xmemory[i] = x[i]
    styStore Ymemory[i] = y[i]
    stz / stziStore Zmemory[i] = z[_][i] (Interleaved Z)

  3. Understand the `vecfp` instruction

    main

    The vecfp instruction performs a pointwise fused-multiply-add or other ALU operation between an X vector, a Y vector, and a Z vector, accumulating the result onto the Z vector. The general mathematical form is:

    z[_][i] ±= f(x[i], y[i])

    Key characteristics:

    • Accumulation: Results are accumulated into the Z vector.
    • Data Types: All three vectors typically share the same element type (f16, f32, or f64). On M2, bf16 is supported.
    • Mixed Lane Widths: When X and Y are f16 (or bf16 on M2), Z can be f32, utilizing two rows of Z for interleaved pairs.
    • M2 Extensions: On M2 hardware, the operation can be repeated multiple times (2 or 4 iterations) using bit 31 and bit 25 of the operand.

  4. Understand the genlut instruction modes

    main

    The genlut instruction operates in two primary modes based on the mode bits (bits 53-56 of the operand):

    1. Generate Mode (Modes 0-6): Used to generate indices for indexed loads. It reads a DT[LC] vector from a source and searches a table register. For each lane in the source, it finds the minimum v such that table[v] > source_lane, and the resulting index is v - 1. If no such v is found, the index is -1 (represented as an unsigned integer with all bits set, except for f64 where the high bit is forced to zero). The resulting indices are written to the low bytes of an X or Y register.

    2. Lookup Mode (Modes 7-15): Performs an indexed load. It reads a densely-packed IT[LC] vector from the source, treats each lane as an index into a table register (of type DT[LC]), expands the lanes to the destination data type DT, and writes the 64-byte result to an X, Y, or Z register.

    Usage Pattern: A common pattern for approximating unary functions is to use a Generate mode to determine which piece of a function a value falls into, followed by a Lookup mode to fetch the relevant coefficients for that piece.

  5. Understand performance benchmarks for vecfp on M1 Max

    main

    The vecfp performance measurements on M1 Max hardware are based on fused-multiply-add (FMA) operations.

    Key considerations for interpreting benchmarks:

    • FMA Counting: A single fused-multiply-add counts as two floating-point operations.
    • GFLOPS Definition: 1.0 GFLOPS equals $10^9$ floating-point operations per second.
    • Workload Type: The provided measurements are performed without any load or store instructions. Real-world workloads involving memory access (loads/stores) will achieve lower performance numbers than those listed.
    • ALU Operations: The benchmarks assume ALU operations of z + x*y or z - x*y.

  6. Understand `matint` data layouts for ALU mode 8

    main

    When alumode is 8, FMA is performed with 8-bit multiplicands, accumulating onto 16-bit or 32-bit $Z$, using a specific data layout.

    32-bit Z Layout

    Each 4x4 block of bytes follows this pattern:

    $X_0$$X_1$$X_2$$X_3$
    $Y_0$$\text{colspan 4}$$Z_{0,0:3} += X_0 \times Y_0$
    $Y_1$$\text{colspan 4}$$Z_{1,0:3} += X_1 \times Y_0$
    $Y_2$$\text{colspan 4}$$Z_{2,0:3} += X_2 \times Y_0$
    $Y_3$$\text{colspan 4}$$Z_{3,0:3} += X_3 \times Y_0$

    16-bit Z Layout

    Each 2x2 block of bytes follows this pattern:

    $X_0$$X_1$
    $Y_0$$\text{colspan 2}$$Z_{0,0:1} += X_0 \times Y_0$
    $Y_1$$\text{colspan 2}$$Z_{1,0:1} += X_1 \times Y_0$

    M3 Specific: 32-bit Z with 16-bit Y

    On M3 machines, alumode 8 also supports 8-bit $X$, 16-bit $Y$, and 32-bit $Z$:

    $X_0$$X_1$$X_2$$X_3$
    $Y_0$$\text{colspan 4}$$Z_{0,0:3} += X_0 \times Y_{0:1}$
    $Y_1$$\text{colspan 4}$$Z_{1,0:3} += X_1 \times Y_{0:1}$
    $Y_2$$\text{colspan 4}$$Z_{2,0:3} += X_2 \times Y_{0:1}$
    $Y_3$$\text{colspan 4}$$Z_{3,0:3} += X_3 \times Y_{0:1}$

  7. Understand the `matfp` instruction

    main

    The matfp instruction performs a fused-multiply-add (or other ALU operation) outer-product between an X vector, a Y vector, and a 2D grid of Z values, accumulating the result onto Z.

    Key Characteristics:

    • Operation: z[j][i] ±= f(x[i], y[j])
    • Data Types: X, Y, and Z share the same element type (f16, f32, or f64). On M2 machines, bf16 is also supported.
    • Mixed Precision: When X and Y are f16 (or bf16 on M2), Z can be f32. In this mode, the entire 64x64 byte grid of Z is used, with even lanes of X mapping to even Z registers and odd lanes of X mapping to odd Z registers.
    • Features: Supports indexed loads, X/Y shuffles, and positive selection (masking).

  8. Understand the `extrh` instruction overview

    main

    The extrh instruction is used for extracting data from Z registers into X or Y registers. Its behavior depends on the value of bit 26 in the instruction encoding:

    • When bit 26 = 0: Performs a simple copy where x[i] = z[_][i]. The lane width of X, Y, and Z must be the same. The operation uses a 7-bit writemask.
    • When bit 26 = 1: Performs a narrowing operation where x[i] = f(z[_][i]) or y[i] = f(z[_][i]). This supports mixed lane widths (e.g., extracting 16-bit elements from 32-bit Z elements). It uses a 9-bit writemask and supports optional integer right shifts and integer saturation.

  9. Understand the `matint` instruction types

    main

    The matint instruction operates in two primary modes depending on the value of bits 47:4 (the ALU mode field).

    1. Reduction Mode (47=4): Performs an in-place reduction of a 2D grid of $Z$ values. This involves a right shift (either rounding or truncating), optionally followed by saturation to various integer widths ($i8, u8, i16, u16, i32, u32$). $Z$ values are 16-bit or 32-bit integers.

    2. Outer-Product Mode (47≠4): Performs an ALU operation in an outer-product manner between an $X$ vector, a $Y$ vector, and a 2D grid of $Z$ values, accumulating the result onto $Z$. This mode supports various combinations of lane widths and special features like shuffles, indexed loads, and popcnt operations.

  10. Use indexed loads in AMX

    main

    While standard loads operate on a 64-byte span, some operations support indexed loads.

    An indexed load is parameterized by:

    • Element Size (ES): 8, 16, 32, or 64 bits.
    • Index Size (IS): 2, 4, or 5 bits.

    Instead of loading a full 512-bit span, an indexed load loads an IS * EC bit span (where EC is the element count) and treats every group of IS bits as a lane index into a different register of size ES.

    Example: With ES=16 (f16) and IS=2, a 64-bit span is loaded from X or Y, viewed as a u2[32] vector, and used to look up lanes in an f16[32] vector.

  11. Hardware variations across Apple Silicon generations

    main

    AMX instruction support varies depending on the chip generation. When developing for AMX, be aware of these historical changes:

    • M1: Likely uses AMX version 2 (supports 7-bit or 9-bit writemasks).
    • M2: Adds bf16 support and other tweaks.
    • M3: Adds one extra mode to each of ldx, ldy, and matint.
    • M4: Certain modes of extrh, extrv, vecfp, and vecint ignore the low bits of the X/Y offset.

  12. Understand the `vecint` instruction

    main

    The vecint instruction is a versatile integer vector operation used on Apple silicon (M1, M2, M3, M4). It performs two primary types of operations depending on the value of bit 47:

    1. In-place Reduction (when 47=4): Performs a right shift (rounding or truncating) on an integer vector in register Z, with an optional saturation step to various integer widths (i8/u8/i16/u16/i32/u32).
    2. ALU Accumulation (when 47≠4): Performs an ALU operation between vectors X, Y, and Z, accumulating the result back into Z. The general form is z[_][i] ±= f(x[i], y[i]).

    Key features include indexed loads, shuffles, broadcast modes, and (on M2+) multi-vector repetition.