Intel Advanced Matrix Extensions (AMX) provide high-throughput matrix operations through architected two-dimensional tile registers. Correct use of these instructions, however, depends on an input layout constraint that is not apparent from the conventional matrix-multiplication formulation: one operand must be stored in Vector Neural Network Instructions (VNNI) format. This note develops that constraint from the AMX data path, demonstrates its effect on correctness, and relates the low-level interface to LIBXSMM’s JIT-dispatch model.
The layout is named after the AVX-512 VNNI dot-product instructions for which it was introduced; AMX inherits the same packing convention.
AMX execution model
Conventional vector registers represent one-dimensional sequences. In contrast, an AMX tile register represents a two-dimensional byte array.
AVX-512 register (1D) AMX tile register (2D)
┌────────────────────┐ ┌────────────────────┐
│ 3 7 1 4 9 2 …│ │ 3 7 1 4 … │
└────────────────────┘ │ 9 2 5 8 … │
one row of 16 numbers │ 6 0 3 1 … │
│ ⋮ │
└────────────────────┘
16 rows × 64 bytes
The architecture provides eight tile registers, tmm0 through tmm7, and a tile matrix-multiply unit (TMUL). A TMUL instruction consumes two source tiles and accumulates their product into a destination tile.
The remainder of the interface configures tile dimensions, moves tiles between memory and the register file, and arranges operands in the format expected by the TMUL datapath.
Matrix multiplication at tile granularity
Consider the product of and :
A B
┌ ┐ ┌ ┐
│ 1 2 3 │ │ 1 0 │
│ 4 5 6 │ │ 0 1 │
└ ┘ │ 2 2 │
└ ┘
The entry C[0][0] is the dot product of row 0 of A and column 0 of B:
row0 of A: 1 2 3
× × ×
col0 of B: 1 0 2
↓ ↓ ↓
1 + 0 + 6 = 7
The resulting matrix is:
C
┌ ┐
│ 7 8 │
│ 16 17 │
└ ┘
This computation performs twelve scalar multiplications. It can be mapped to the machine at several granularities:
- Scalar CPU: execute twelve scalar multiply instructions.
- AVX-512: load
[1,2,3]and[1,0,2], multiply, and horizontally reduce. This approach evaluates one dot product per vector operation sequence. - AMX: load A and B into tiles and issue a single TMUL instruction. The TMUL array evaluates all four output dot products concurrently.
┌──────────┐ ┌──────────┐
│ A │ │ B │
│ (tmm1) │ │ (tmm2) │
└────┬─────┘ └────┬─────┘
│ │
└──────► TMUL ◄──────┘ one instruction
│
▼
┌──────────┐
│ C │
│ (tmm0) │
└──────────┘
The TMUL array is organized around the output tile: each cell evaluates one dot product independently. Instruction latency is fixed with respect to active cells; partially filled tiles therefore reduce effective utilization.
AMX BF16 tile dimensions
For BF16 inputs, the relevant AMX tile operation has the following dimensions:
A B C
16 × 32 × 32 × 16 = 16 × 16
(bf16) (bf16) (fp32)
The operation produces 256 output elements, each a length-32 dot product. It therefore executes 8,192 multiply-adds in one instruction, with an approximate latency of 16 cycles, or 1,024 FLOPs per cycle per core.
An AVX-512 FMA sustains 32 FLOPs per cycle. The resulting 32-fold arithmetic-throughput difference motivates AMX for dense low-precision matrix multiplication.
FLOP accounting
For an output element, a fresh dot product appears to require K multiplications and K−1 additions, or 2K−1 operations. Nevertheless, standard throughput accounting reports 2K FLOPs. The distinction follows from the instruction semantics.
The instruction implements tmm0 += tmm1 × tmm2; the destination is an accumulator across the K loop rather than a fresh assignment:
C[0][0] = 0 ← the accumulator exists before you start
C[0][0] += 1·1 ← mult, add
C[0][0] += 2·0 ← mult, add
C[0][0] += 3·2 ← mult, add
3 mults + 3 adds = 2K
The additional addition arises from the accumulation operation and corresponds to a hardware FMA. Thus, ; for one BF16 tile operation, this is .
Utilization constraint: reduction dimension
The 3 in the matrix A is the dot-product, or reduction, dimension. AMX fixes this dimension at 32 for BF16 and 64 for INT8 operations.
K = 32 → every multiply unit busy 100% of peak
K = 16 → half the grid computes nothing 50%
K = 4 → most of the grid idles 12.5%
The instruction retains approximately the same latency in each case. A short reduction dimension activates only a subset of the TMUL array and proportionally reduces utilization. This fixed cost is a common reason that manually written AMX kernels fail to reach expected performance.
For , the kernel issues the tile instruction 128 times against the same C tile:
C_tile ← 0
for k_chunk in 0..127:
load A_chunk, B_chunk
C_tile += A_chunk × B_chunk
store C_tile ← once, at the very end
C remains resident in the tile register file while A and B stream from memory. This reuse amortizes output traffic and explains the natural accumulate semantics.
VNNI operand layout
The key layout constraint follows directly from the TMUL operand access pattern.
TMUL does not consume scalar elements individually. It consumes a 4-byte dword, which contains exactly two BF16 values. It multiplies the corresponding pairs and accumulates their sum, thereby consuming two reduction-dimension elements per fetch.
one fetch from A: ( a[k], a[k+1] )
one fetch from B: ( b[k], b[k+1] )
↓
a[k]·b[k] + a[k+1]·b[k+1]
A dword is a 4-byte memory unit. The hardware operates on this representation without knowledge of high-level matrix semantics.
For A, this constraint is naturally satisfied because consecutive K values are contiguous within a row.
For a conventional row-major B, consecutive K values for a fixed output column reside in different rows. B must therefore be transformed so that each required pair is contiguous.
Deriving the transformation
Consider an even-K example with and :
A B
┌ ┐ ┌ ┐
│ 1 2 3 4 │ │ 1 0 │ k=0
│ 5 6 7 8 │ │ 0 1 │ k=1
└ ┘ │ 2 2 │ k=2
│ 3 1 │ k=3
└ ┘
n=0 n=1
The reference computation is C[0][0] = 1·1 + 2·0 + 3·2 + 4·3 = 19.
The required packing transformation is:
B_packed[ k/2 ][ 2n + (k mod 2) ] = B[k][n]
Applying the transformation to each B element yields:
B[0][0]=1 → row 0, col 0 B[0][1]=0 → row 0, col 2
B[1][0]=0 → row 0, col 1 B[1][1]=1 → row 0, col 3
B[2][0]=2 → row 1, col 0 B[2][1]=2 → row 1, col 2
B[3][0]=3 → row 1, col 1 B[3][1]=1 → row 1, col 3
A (unchanged) B (on paper) B_packed (in memory)
┌ ┐ ┌ ┐ ┌─────────┬─────────┐
│ 1 2 3 4 │ │ 1 0 │ k=0 │ 1 0 │ 0 1 │ ← k=0,1
│ 5 6 7 8 │ │ 0 1 │ k=1 ├─────────┼─────────┤
└ ┘ │ 2 2 │ k=2 │ 2 3 │ 2 1 │ ← k=2,3
dword0 dword1 │ 3 1 │ k=3 └─────────┴─────────┘
└ ┘ n=0 n=1
n=0 n=1 each cell = one dword = one K-pair
The dot product then proceeds through two dword fetches:
Fetch 1:
A row 0, dword 0 = (1, 2) ← k = 0,1
B_packed row 0, n=0 dword = (1, 0) ← k = 0,1
1·1 + 2·0 = 1
Fetch 2:
A row 0, dword 1 = (3, 4) ← k = 2,3
B_packed row 1, n=0 dword = (2, 3) ← k = 2,3
3·2 + 4·3 = 18
C[0][0] = 19 ✓
Packing ensures that the second fetch reads (2, 3), the and values of column 0, as an adjacent pair.
Consequence of omitting VNNI packing
If the row-major B representation is merely reinterpreted as two rows of four values:
row 0: [ 1 0 0 1 ]
row 1: [ 2 2 3 1 ]
Right shape. Loads fine. No fault. And then:
Fetch 2:
A row 0, dword 1 = (3, 4)
row 1, n=0 dword = (2, 2) ← B[2][0] and B[2][1] !
two different COLUMNS, same k
3·2 + 4·2 = 14
C[0][0] = 15 ✗ (should be 19)
dwords. The result is incorrect, but no architectural fault is raised. The hardware has multiplied a column-0 element by a column-1 element because the requested dword does not encode the intended matrix column. The TMUL unit only observes dword addresses and contents.
This failure mode can evade naive testing: an all-ones input produces the expected result despite the incorrect layout. Correctness tests should therefore use nonuniform values that distinguish both reduction indices and output columns.
Direct AMX implementation
The following minimal program configures AMX tiles and computes one BF16 tile product directly with intrinsics.
// Build: gcc -O2 -mamx-tile -mamx-bf16 amx_toy.c -o amx_toy
// Needs: Sapphire Rapids or newer, Linux 5.16+
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <immintrin.h>
#include <sys/syscall.h>
#include <unistd.h>
#define ARCH_REQ_XCOMP_PERM 0x1023
#define XFEATURE_XTILEDATA 18
typedef struct {
uint8_t palette_id;
uint8_t start_row;
uint8_t reserved[14];
uint16_t colsb[16]; // bytes per row, per tile
uint8_t rows[16]; // rows, per tile
} __attribute__((packed)) tileconfig;
#define BF16_ONE 0x3F80 // top 16 bits of the fp32 pattern for 1.0f
int main(void) {
// Ask the kernel for permission. Skip this and the first tile
// instruction dies with SIGILL. Nobody warns you about this.
if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) {
fprintf(stderr, "no AMX here\n");
return 1;
}
// tile0 = C: 16x16 fp32. tile1 = A: 16x32 bf16. tile2 = B: 16x32 bf16.
// All three are 16 rows x 64 bytes. The hardware counts bytes, not elements.
tileconfig cfg;
memset(&cfg, 0, sizeof(cfg));
cfg.palette_id = 1;
for (int t = 0; t < 3; t++) { cfg.rows[t] = 16; cfg.colsb[t] = 64; }
_tile_loadconfig(&cfg);
uint16_t A[16][32], B[16][32];
float C[16][16];
for (int m = 0; m < 16; m++)
for (int k = 0; k < 32; k++) A[m][k] = BF16_ONE;
for (int r = 0; r < 16; r++)
for (int c = 0; c < 32; c++) B[r][c] = BF16_ONE;
memset(C, 0, sizeof(C));
_tile_loadd(0, C, 16 * sizeof(float)); // strides are in BYTES
_tile_loadd(1, A, 32 * sizeof(uint16_t));
_tile_loadd(2, B, 32 * sizeof(uint16_t));
_tile_dpbf16ps(0, 1, 2); // C += A * B
_tile_stored(0, C, 16 * sizeof(float));
_tile_release(); // let the core leave the AMX
// power license
printf("C[0][0] = %.1f (expected 32.0)\n", C[0][0]);
return 0;
}
Three constraints are essential:
- The
arch_prctlcall is mandatory. Linux protects tile state through a per-process permission request. Omitting it causes the first tile instruction to terminate withSIGILL. - B must be VNNI-packed. All-ones inputs conceal layout errors; use distinct values when validating the kernel.
- The tile shape is fixed. One
_tile_dpbf16psalways evaluates a 16x32 by 32x16 product. A general GEMM is implemented as a loop nest around this primitive, typically retaining C intmm0..3, A intmm4..5, and B intmm6..7.
LIBXSMM implementation
LIBXSMM is an open-source JIT library developed by Intel Labs (Heinecke et al., SC'16). For selected BF16 workloads, it has been reported to outperform oneDNN by approximately 1.45x, in part because its blocked B layout can avoid cache-conflict misses associated with a leading dimension of 4096.
The following version expresses the same example through LIBXSMM. The API follows the column-major BLAS convention.
/* C(2x2) = A(2x4) * B(4x2), bf16 in, fp32 out.
* LIBXSMM is COLUMN-MAJOR: A[m][k] lives at a[m + k*lda]. */
#include <libxsmm.h>
#include <stdio.h>
int main(void) {
const libxsmm_blasint m = 2, n = 2, k = 4;
float a_f32[8] = { 1,5, 2,6, 3,7, 4,8 }; /* [1 2 3 4; 5 6 7 8] */
float b_f32[8] = { 1,0,2,3, 0,1,2,1 }; /* [1 0; 0 1; 2 2; 3 1] */
libxsmm_bfloat16 A[8], B[8], B_vnni[8];
float C[4] = { 0 };
libxsmm_rne_convert_fp32_bf16(a_f32, A, 8);
libxsmm_rne_convert_fp32_bf16(b_f32, B, 8);
/* --- pack B. The ugly detail, as a library call. --- */
libxsmm_meltw_unary_shape ushape = libxsmm_create_meltw_unary_shape(
n, k, k, k,
LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16);
libxsmm_meltwfunction_unary vnni = libxsmm_dispatch_meltw_unary(
LIBXSMM_MELTW_TYPE_UNARY_TRANSFORM_NORM_TO_VNNI2,
ushape, LIBXSMM_MELTW_FLAG_UNARY_NONE);
libxsmm_meltw_unary_param uparam;
uparam.in.primary = B;
uparam.out.primary = B_vnni;
vnni(&uparam);
/* --- dispatch: JIT a kernel for exactly these sizes --- */
libxsmm_gemm_shape shape = libxsmm_create_gemm_shape(
m, n, k,
m, k, m, /* lda, ldb, ldc */
LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16,
LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32);
libxsmm_bitfield flags = LIBXSMM_GEMM_FLAGS('N','N')
| LIBXSMM_GEMM_FLAG_VNNI_B /* B is pre-packed */
| LIBXSMM_GEMM_FLAG_BETA_0; /* C = A*B, not += */
libxsmm_gemmfunction kernel =
libxsmm_dispatch_gemm(shape, flags, LIBXSMM_GEMM_PREFETCH_NONE);
if (!kernel) { fprintf(stderr, "dispatch failed\n"); return 1; }
/* --- run --- */
libxsmm_gemm_param p;
p.a.primary = A;
p.b.primary = B_vnni;
p.c.primary = C;
kernel(&p);
/* [19 12 ; 43 28], column-major */
printf("C = [ %.0f %.0f ; %.0f %.0f ]\n", C[0], C[2], C[1], C[3]);
return 0;
}
The application contains no tile intrinsics. LIBXSMM internally manages arch_prctl, tileconfig, and _tile_release; the same source can target AVX-512, AVX2, and ARM SVE without architecture-specific conditionals or compiler flags.
the flag, and the same silent-wrong-answer bug reappears.
LIBXSMM_GEMM_FLAG_VNNI_B is an input contract, rather than a request to perform packing. It informs the generated kernel that B is already packed. Setting the flag while omitting the transform recreates the same silent correctness failure described above.
JIT dispatch semantics
LIBXSMM uses a two-phase dispatch-and-execute interface.
An ordinary general-purpose GEMM implementation performs inner-loop bookkeeping: it must identify the final iteration, determine whether K is divisible by the tile depth, and execute remainder paths where necessary. These branches add overhead, while a single binary must support both and .
LIBXSMM instead specializes generated code to the known , , and dimensions before execution.
You: "m=2, n=2, k=4, bf16 in, fp32 out, B pre-packed."
│
[ D I S P A T C H ]
│
emits x86 bytes into an executable page:
tileloadd tmm1, [rdi]
tileloadd tmm2, [rsi]
tdpbf16ps tmm0, tmm1, tmm2
tilestored [rdx], tmm0
ret
│
function pointer
▼
You: kernel(¶m);
The resulting specialized kernel can eliminate loop counters and branches. The relevant loop is unrolled during JIT generation because the trip count is then constant, even when it was not known during application compilation.
| cost | how often | |
|---|---|---|
| dispatch | microseconds | once, at startup |
| run | nanoseconds | millions of times |
Dispatch should occur outside the hot loop. Although LIBXSMM caches kernels by shape, the lookup overhead can exceed the execution time of a small specialized kernel.
The VNNI transform follows the same shape, dispatch, and execute sequence. In LIBXSMM terminology, meltw denotes an elementwise operation, unary denotes one input and one output tensor, and NORM_TO_VNNI2 selects the relevant layout transformation from an operation family that also includes ReLU and transpose.
In inference deployments, the transform should generally not occur on the critical path. B commonly represents model weights, which can be packed once at model load time and retained in the packed representation. A represents per-token activations and requires no equivalent transformation; only B requires the VNNI shuffle.
Precision and bottleneck analysis
Lower precision does not universally imply higher AMX compute throughput.
unpacked to INT8 before the tile op.
AMX provides no native INT4 tile instruction. The ISA defines tdpbf16ps for BF16, the tdpb{ss,su,us,uu}d family for INT8, and tdpfp16ps for FP16 on Granite Rapids. Any AMX INT4 execution path uses 4-bit storage and unpacks values to INT8 before the tile operation.
| elem/row | K per op | ops/cycle/core | |
|---|---|---|---|
| BF16 | 32 | 32 | 1,024 FLOP |
| INT8 | 64 | 64 | 2,048 OP |
| INT4 | — | — | 2,048 OP (via INT8) |
INT8 doubles the BF16 operation count because a 64-byte row contains twice as many elements. INT4 provides no additional tile compute throughput and introduces unpacking overhead.
The relevant choice depends on whether inference is compute-bound or bandwidth-bound.
PREFILL — compute-bound. Weights reused across many tokens.
BF16 ████████████████ ~288 TFLOPS
INT8 ████████████████████████████████ ~576 TOPS (2x)
INT4 ██████████████████████████████ ~576 minus unpack
DECODE — bandwidth-bound. Each weight read once, used once.
bytes/param → arithmetic intensity → ceiling
BF16 2.0 B 1 OP/B ████ ~0.85 TFLOPS
INT8 1.0 B 2 OP/B ████████ ~1.7 TOPS (2x)
INT4 0.5 B 4 OP/B ████████████████ ~3.4 TOPS (4x)
During decode, performance is dominated by bytes transferred rather than operations issued. INT4 halves weight traffic, while the TMUL unit is largely idle awaiting DRAM; unpacking can therefore be effectively hidden. During prefill, compute is the primary bottleneck. INT4 provides no additional AMX throughput, and its shift-and-mask unpack sequence reduces the potential benefit.
Capacity is a separate consideration. For a 671B-parameter model such as DeepSeek-V3:
| format | footprint |
|---|---|
| BF16 | ~1.34 TB |
| INT8 | ~671 GB |
| INT4 | ~336 GB |
INT4 can make single-node inference feasible at this scale. Its principal benefit in this setting is capacity, rather than raw AMX compute throughput.
Numeric storage format and arithmetic execution format are independent design choices.
Portions of this post were drafted with the assistance of AI and subsequently reviewed and edited for technical accuracy.