mac16.md

November 16, 2022 · View on GitHub

Quick summary

InstructionGeneral themeWritemaskOptional special features
mac16 (63=0)z[j][i] += x[i] * y[j]7 bit X, 7 bit YX/Y/Z input disable, right shift
mac16 (63=1)z[_][i] += x[i] * y[i]7 bitX/Y/Z input disable, right shift

Instruction encoding

BitWidthMeaningNotes
1022A64 reserved instructionMust be 0x201000 >> 10
55InstructionMust be 14
055-bit GPR indexSee below for the meaning of the 64 bits in the GPR

Operand bitfields

BitWidthMeaningNotes
631Vector mode (1) or matrix mode (0)
621Z is i32 (1) or Z is i16 (0)Ignored in vector mode; Z is always i16 there
611X is i8 (1) or X is i16 (0)
601Y is i8 (1) or Y is i16 (0)
555Right shift amountApplied to x*y. When zero, sign of x and y inputs is less relevant.
487Ignored
462X enable mode
415X enable valueMeaning dependent upon associated mode
392Ignored
372Y enable modeIgnored in vector mode
325Y enable valueIgnored in vector mode
Meaning dependent upon associated mode
302Ignored
291Skip X input (1) or use X input (0)
281Skip Y input (1) or use Y input (0)
271Skip Z input (1) or use Z input (0)
261Ignored
206Z rowHigh bits ignored in matrix mode
191Ignored
109X offset (in bytes)
91Ignored
09Y offset (in bytes)

Combinations of bits 27-29 result in various integer ALU operations:

Operation29 (X)28 (Y)27 (Z)
z+((x*y)>>s)000
   (x*y)>>s 001
z+( x   >>s)010
    x   >>s 011
z+(   y >>s)100
      y >>s 101
z           110
0111

Combinations of bits 60-63 result in various widths / signs for X / Y / Z:

ModeXYZ63 (M)62 (Z)61 (X)60 (Y)
Matrixi16 [a]i16 [a]i16 or u16 (one row from each two)0000
Matrixi16 [a]i8 (even lanes)i16 or u16 (one row from each two)0001
Matrixi8 (even lanes)i16 [a]i16 or u16 (one row from each two)0010
Matrixi8 (even lanes)i8 (even lanes)i16 or u16 (one row from each two)0011
Matrixi16i16i32 or u32 (all rows, interleaved pairs)0100
Matrixi16i8 (even lanes)i32 or u32 (all rows, interleaved pairs)0101
Matrixi8 (even lanes)i16i32 or u32 (all rows, interleaved pairs)0110
Matrixi8 (even lanes)i8 (even lanes)i32 or u32 (all rows, interleaved pairs)0111
Vectori16 [a]i16 [a]i16 or u16 (one row)100
Vectori16 [a]i8 (even lanes)i16 or u16 (one row)101
Vectori8 (even lanes)i16 [a]i16 or u16 (one row)110
Vectori8 (even lanes)i8 (even lanes)i16 or u16 (one row)111

[a] Or u16 if shift amount is zero.

X/Y enable modes:

ModeMeaning of value (N)
0Enable all lanes (0), or odd lanes only (1), or even lanes only (2), or no lanes (anything else)
1Only enable lane #N
2Only enable the first N lanes, or all lanes when N is zero
3Only enable the last N lanes, or all lanes when N is zero

Description

In vector mode, takes an X vector of type i8 or i16, a Y vector of type i8 or i16, and a Z vector of type i16, and performs pointwise: multiply X by Y, right shift (truncating) by some amount, then add on to Z. Variants of this pointwise operation remove the X and/or Y and/or Z inputs. When X or Y have type i8, the 8 bits are taken from the low 8 bits of each 16-bit lane.

In matrix mode, takes an X vector of type i8 or i16, a Y vector of type i8 or i16, and a 2D grid of Z values of type i16 or i32, and performs an outer product of X and Y followed by pointwise right shift (truncating) by some amount, and then pointwise addition onto Z. Variants of this remove the X and/or Y and/or Z inputs. When X or Y have type i8, the 8 bits are taken from the low 8 bits of each 16-bit lane. When Z has type i32, the entire 64x64 byte grid of Z is used, with even lanes of X going into even Z registers and odd lanes of X going into odd Z registers (see Mixed lane widths).

Emulation code

See mac16.c.

A representative sample is:

void emulate_AMX_MAC16(amx_state* state, uint64_t operand) {
    uint64_t y_offset = operand & 0x1FF;
    uint64_t x_offset = (operand >> 10) & 0x1FF;
    uint64_t z_row = (operand >> 20) & 63;
    uint64_t x_enable = parse_writemask(operand >> 41, 2, 7);
    uint64_t y_enable = parse_writemask(operand >> 32, 2, 7);

    int16_t x[32];
    int16_t y[32];
    load_xy_reg(x, state->x, x_offset);
    load_xy_reg(y, state->y, y_offset);

    for (int i = 0; i < 32; i++) {
        if (!((x_enable >> (i * 2)) & 1)) continue;
        if (operand & FMA_VECTOR_PRODUCT) {
            int16_t* z = &state->z[z_row].i16[i];
            *z = mac32_alu(x[i], y[i], *z, operand);
        } else {
            for (int j = 0; j < 32; j++) {
                if (!((y_enable >> (j * 2)) & 1)) continue;
                if (operand & FMA_WIDEN_16_32) {
                    int32_t* z = &state->z[(j * 2) + (i & 1)].i32[i >> 1];
                    *z = mac32_alu(x[i], y[j], *z, operand);
                } else {
                    int16_t* z = &state->z[(j * 2) + (z_row & 1)].i16[i];
                    *z = mac32_alu(x[i], y[j], *z, operand);
                }
            }
        }
    }
}

int64_t mac32_alu(int64_t x, int64_t y, int64_t z, uint64_t operand) {
    if (operand & MAC_X_INT8) x = (int8_t)x;
    if (operand & MAC_Y_INT8) y = (int8_t)y;
    int64_t val;
    switch ((operand >> 28) & 3) {
    default: val = x * y; break;
    case  1: val = x; break;
    case  2: val = y; break;
    case  3: val = 0; break;
    }
    uint32_t shift = (operand >> 55) & 0x1f;
    val >>= shift;
    if (!(operand & MAC_SKIP_Z_INPUT)) {
        val += z;
    }
    return val;
}

Performance (M1 Max)

Note that multiply-and-add counts as two operations. A measurement of 1.0 GOPS would mean 109 operations per second. The measurements are done without any load or store instructions; real-world workloads will need loads and stores, and thus will achieve lower numbers.

mac16 in matrix mode, with both of X and Y being i8, and each Z accumulator being i16[32][32]:

Z Accumulators1 Thread2 Threads3 Threads4 Threads5 Threads6 Threads
1 per thread1450.7 GOPS2885.9 GOPS2697.9 GOPS3566.1 GOPS4434.0 GOPS5199.1 GOPS
2 per thread2964.4 GOPS5821.3 GOPS4866.4 GOPS6196.0 GOPS5601.5 GOPS6286.5 GOPS

mac16 in matrix mode, with X or Y or both being i16, and each Z accumulator being i16[32][32]:

Z Accumulators1 Thread2 Threads3 Threads4 Threads5 Threads6 Threads
1 per thread1459.3 GOPS2290.9 GOPS2634.3 GOPS2370.3 GOPS2863.8 GOPS2875.9 GOPS
2 per thread1467.4 GOPS2362.0 GOPS2556.9 GOPS2303.7 GOPS2790.0 GOPS2907.3 GOPS

mac16 in matrix mode, with both of X and Y being i8, and each Z accumulator being i32[32][32]:

Z Accumulators1 Thread2 Threads3 Threads4 Threads5 Threads6 Threads
1 per thread1476.7 GOPS2908.5 GOPS2647.6 GOPS2184.1 GOPS2806.8 GOPS2892.8 GOPS

mac16 in matrix mode, with X or Y or both being i16, and each Z accumulator being i32[32][32]:

Z Accumulators1 Thread2 Threads3 Threads4 Threads5 Threads6 Threads
1 per thread1454.4 GOPS2947.7 GOPS2198.5 GOPS2397.4 GOPS2730.1 GOPS2844.9 GOPS

mac16 in vector mode, with both of X and Y being i8, and each Z accumulator being i16[32]:

Z Accumulators1 Thread2 Threads3 Threads4 Threads5 Threads6 Threads
1 per thread45.6 GOPS92.6 GOPS109.6 GOPS129.3 GOPS162.1 GOPS195.4 GOPS
2 per thread91.4 GOPS183.9 GOPS204.6 GOPS272.9 GOPS385.3 GOPS339.9 GOPS
3 per thread138.3 GOPS277.3 GOPS313.4 GOPS413.1 GOPS473.9 GOPS477.8 GOPS
4 per thread180.4 GOPS363.6 GOPS416.9 GOPS545.6 GOPS635.5 GOPS610.5 GOPS
5 per thread231.0 GOPS461.2 GOPS482.1 GOPS642.7 GOPS751.2 GOPS648.2 GOPS
6 per thread278.0 GOPS551.4 GOPS543.4 GOPS722.0 GOPS771.1 GOPS759.5 GOPS
7 per thread318.5 GOPS644.5 GOPS598.9 GOPS754.3 GOPS785.0 GOPS792.5 GOPS
8 per thread369.8 GOPS735.9 GOPS669.0 GOPS799.9 GOPS789.6 GOPS757.2 GOPS

mac16 in vector mode, with X or Y or both being i16, and each Z accumulator being i16[32]:

Z Accumulators1 Thread2 Threads3 Threads4 Threads5 Threads6 Threads
1 per thread46.0 GOPS92.6 GOPS107.7 GOPS142.1 GOPS176.9 GOPS199.9 GOPS
2 per thread92.5 GOPS185.2 GOPS211.4 GOPS246.3 GOPS281.9 GOPS287.1 GOPS
3 per thread137.7 GOPS277.0 GOPS321.5 GOPS388.3 GOPS463.6 GOPS415.3 GOPS
4 per thread185.5 GOPS369.4 GOPS425.8 GOPS461.4 GOPS486.6 GOPS481.3 GOPS
5 per thread173.9 GOPS346.1 GOPS379.1 GOPS426.7 GOPS463.8 GOPS449.3 GOPS
6 per thread185.2 GOPS361.6 GOPS427.5 GOPS455.2 GOPS476.7 GOPS465.4 GOPS
7 per thread185.6 GOPS368.7 GOPS400.8 GOPS444.5 GOPS469.7 GOPS455.6 GOPS
8 per thread185.1 GOPS369.2 GOPS413.8 GOPS464.1 GOPS477.5 GOPS465.6 GOPS