Intel AMX in C/C++: Registers, Intrinsics, and Data Flow

Intel AMX (Advanced Matrix Extensions) adds matrix instructions to x86 processors. It introduces two-dimensional tile registers and instructions that accumulate the product of two matrix blocks into a third. Fully connected neural network layers, convolutions expressed as matrix multiplication, and batches of vector dot products can all organize their inner computations this way.¹

AMX first appeared in 4th Gen Xeon Scalable processors (Sapphire Rapids), with INT8 and BF16 support. The 5th Gen Xeon Scalable family (Emerald Rapids) retains both. Xeon 6 processors with P-cores (Granite Rapids) also support AMX and add FP16; the E-core products based on Sierra Forest do not support AMX. Some workstation processors support it as well, including the w9-3495X in the Xeon W-3400 series. The Xeon Platinum 8458P used here belongs to Sapphire Rapids. The INT8 examples below target processors with AMX-TILE and AMX-INT8.²³⁴

C/C++ programs access these instructions through intrinsics. To compute C=AB for matrices in memory, a program loads blocks of A and B into tiles, performs multiply-accumulate operations, and stores C back to an array. Loads and stores access ordinary process memory through the CPU cache hierarchy.

The examples start with UINT8 × INT8 → INT32 arithmetic and a fixed 16×16 output block, with correctness and performance tested on a Xeon Platinum 8458P. Each snippet follows a load, computation, or store involving A, B, and C, showing where the data resides in memory and in the tiles.

Headers and compiler options

AMX intrinsics are available through immintrin.h, the same header used for AVX and AVX-512:

1

#include <immintrin.h>

An INT8 kernel uses two instruction extensions: AMX-TILE configures tiles and moves their data, while AMX-INT8 performs integer dot products. Enable both in the compiler. With a GCC toolchain that supports AMX, use:⁵

1

g++ -O3 -std=c++17 -mamx-tile -mamx-int8 amx_int8.cpp -o amx_int8

For BF16 inputs with FP32 accumulation, also enable -mamx-bf16. Compiler flags permit instruction generation; the CPU running the program must implement those extensions, and Linux must allow the process to use tile state.

You can also enable AMX for just the compute function. The target attribute below enables AMX in that function while main and the feature checks use the ordinary x86-64 target. The complete examples later in this article use this approach.

1
2
3
4
5

__attribute__((target("amx-tile,amx-int8"), noinline))
void amx_block(const unsigned char* a, const signed char* packed_b,
               int* c) {
    // AMX intrinsics go here.
}

GCC and Clang both provide the interfaces with explicit tile indices used below. Their definitions are available in the GCC and LLVM headers. The build commands here use GCC. If an intrinsic, target option, or assembly instruction is unrecognized, check both the compiler and assembler versions.⁶⁷

On Linux, request permission from the kernel before the first use of tile data:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

#include <asm/prctl.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <cstdio>

#ifndef ARCH_REQ_XCOMP_PERM
#define ARCH_REQ_XCOMP_PERM 0x1023
#endif

bool enable_amx() {
    constexpr unsigned long XFEATURE_XTILEDATA = 18;
    if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM,
                XFEATURE_XTILEDATA) != 0) {
        std::perror("ARCH_REQ_XCOMP_PERM");
        return false;
    }
    return true;
}

The value 18 is the extended-state component ID. The kernel grants the process permission to use tile data and allocates storage for saving that state when the executing thread first uses it. The permission applies within the process, is inherited across fork, and is cleared by exec. Applications that install an alternate signal stack must also provide enough space for the larger signal frame. The complete example checks the CPU’s AMX-TILE, AMX-INT8, and OSXSAVE features before requesting permission.⁸

Registers and core intrinsics

AMX stores tile shapes in TILECFG and computation data in tmm0 through tmm7. Palette 1, used here, provides 8 tiles, each with up to 16 rows of 64 bytes, for a total capacity of 8 KiB.¹

The rows field specifies the row count, and colsb specifies the number of bytes per row. Setting them to 16 and 64 gives a 1024-byte tile. Each row then holds 64 INT8 elements, 32 BF16 elements, or 16 INT32 elements. The configuration can be identical in all three cases; the compute instruction determines how to interpret the bytes.

The following interfaces configure tiles, compute results, and move data. The arguments dst, a, and b are tile indices; base is a memory address; stride is the distance in bytes between the starts of adjacent memory rows.

These interfaces name registers directly. For example, _tile_zero(0) clears tmm0 and _tile_zero(1) clears tmm1. Tile indices must be known at compile time. Kernels therefore tend to use a fixed set of tiles and change the load addresses inside a loop, reusing those registers for successive matrix blocks.⁶

To compute $C_{M\times N} += A_{M\times K}B_{K\times N}$, assign tmm0 to C, tmm1 to A, and tmm2 to packed B. The INT8 example uses M=16, N=16, and K=64. All three tiles have 16 rows of 64 bytes, but their element types and mathematical roles differ.

Loading: tile shapes and memory layouts

Configure the tile shapes before loading data. The structure below matches the configuration layout: palette_id is at byte 0, colsb starts at offset 16, and rows starts at offset 48. Unused fields remain zero.¹

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

struct alignas(64) TileConfig {
    std::uint8_t palette_id;
    std::uint8_t start_row;
    std::uint8_t reserved0[14];
    std::uint16_t colsb[8];
    std::uint8_t reserved1[16];
    std::uint8_t rows[8];
    std::uint8_t reserved2[8];
};
static_assert(sizeof(TileConfig) == 64);
static_assert(offsetof(TileConfig, rows) == 48);

TileConfig cfg{};
cfg.palette_id = 1;
cfg.rows[0] = 16;  cfg.colsb[0] = 16 * sizeof(std::int32_t); // C
cfg.rows[1] = 16;  cfg.colsb[1] = 64;                       // A
cfg.rows[2] = 16;  cfg.colsb[2] = 16 * 4;                   // packed B
_tile_loadconfig(&cfg);

Load A[16][64] into tmm1. A contains UINT8 elements, with 64 bytes per row and no gaps between rows. The configuration already specifies 16 rows of 64 bytes, so the call only needs the starting address and a row stride of 64 bytes:

1

_tile_loadd(1, &a[0][0], 64);

When loading a block from a larger matrix, colsb and stride often differ. Suppose each row of A contains 128 UINT8 elements, and a 16×64 block starts at row m0 and column k0. Set colsb to 64 and stride to 128:

1
2

// large_a has 128 bytes per row; the configured tile reads 64 bytes per row.
_tile_loadd(1, &large_a[m0][k0], 128);

This load reads 64 bytes from each row, then advances to the next row’s start, 128 bytes away. The other 64 bytes in each row do not enter the tile. After all 16 rows have been read, tmm1 still contains a compact 16×64 bytes; the gaps between memory rows are not stored in the register.

Next, load B. Computing C[m][n] requires the dot product of row m of A and column n of B, but a row-major B stores elements from different columns next to each other. AMX INT8 instructions multiply and accumulate groups of 4 byte pairs, so 4 consecutive K elements from the same column of B must be adjacent. This loop performs the rearrangement:

1
2
3

for (int k = 0; k < K; ++k)
    for (int n = 0; n < N; ++n)
        packed_b[k / 4][4 * n + k % 4] = b[k][n];

The figure uses K=8 and N=4. The first packed row starts with b00, b10, b20, and b30 from column 0, followed by b01, b11, b21, and b31 from column 1, then the four elements from each of columns 2 and 3. The next row handles the remaining 4 elements along K. Colors identify output columns, while the subscripts show each byte’s original matrix coordinates.

For K=64 and N=16, packed B has K/4=16 rows of 4N=64 bytes. It can be loaded directly into tmm2:

1

_tile_loadd(2, &packed_b[0][0], 64);

A tile load copies bytes row by row, so B must be rearranged before loading. Loading the original B directly leaves bytes from different output columns next to each other. The compute instruction would treat them as 4 K elements from the same column. This incorrect pairing can change the result without causing a runtime exception.⁷

Computing: dot products and accumulators

Once the data is loaded, specify the three tiles for the computation:

1
2

_tile_zero(0);
_tile_dpbusd(0, 1, 2);

The first statement clears the accumulator in tmm0. The second reads tmm1 and tmm2 and adds their dot products to tmm0; both input tiles retain their contents. For output position (m,n), the operands are row m of A and the byte groups belonging to column n of packed B. The operation expands into the following scalar loops:

1
2
3
4
5
6

for (int m = 0; m < M; ++m)
    for (int n = 0; n < N; ++n)
        for (int g = 0; g < K / 4; ++g)
            for (int j = 0; j < 4; ++j)
                c[m][n] += static_cast<std::int32_t>(a[m][4 * g + j]) *
                           static_cast<std::int32_t>(packed_b[g][4 * n + j]);

For example, let the first 4 elements of A be [1, 2, 3, 4]. In packed B, output column 0 contains [−1, 2, −3, 4], and column 1 contains [−2, 1, 0, 3]. These groups contribute the following values to the two outputs:

1
2

C[m][0] += 1*(-1) + 2*2 + 3*(-3) + 4*4 = 10
C[m][1] += 1*(-2) + 2*1 + 3*0    + 4*3 = 12

Subsequent K groups continue accumulating into these positions. A tile dot-product instruction performs the same operation for every row and output column covered by the configuration.

Each group multiplies 4 pairs of INT8 values and accumulates them into the corresponding INT32 element. Use dpbusd for UINT8 A and INT8 B, and dpbssd when both inputs are INT8. The byte 0xff represents 255 or −1 under these two interpretations, so test inputs should include unsigned values above 127 and negative signed values.⁷

Each dot product adds to the value already in tmm0. To compute C=AB, clear it before one or more dot products. To compute C+=AB, first load the existing C into tmm0:

1
2

_tile_loadd(0, &c[0][0], N * sizeof(std::int32_t));
_tile_dpbusd(0, 1, 2);

When K is too large for one tile, keep C in tmm0 while loading successive K blocks of A and B. This snippet assumes that A has 128 columns and B has been packed into two separate 64×16 panels, each stored as [16][64]:

1
2
3
4
5
6
7

_tile_zero(0);
for (int block = 0; block < 2; ++block) {
    _tile_loadd(1, &a128[0][block * 64], 128);
    _tile_loadd(2, &b_panels[block][0][0], 64);
    _tile_dpbusd(0, 1, 2);
}
// tmm0 now contains the sum over all 128 K-elements.

Load the configuration outside the accumulation loop; inside it, only load inputs and compute dot products. Reloading the configuration initializes tile data state, so keep the same configuration throughout accumulation.⁹ For multiple output blocks, assign a C tile to each block and reuse the A or B tiles already loaded. Input and output blocks share the same 8 tile registers.

BF16 uses the same data flow, with two BF16 elements per dot-product group and FP32 accumulation. Its intrinsic is:

1

_tile_dpbf16ps(0, 1, 2);

With M=16, N=16, and K=32, A contains 16×32 BF16 elements, C contains 16×16 FP32 elements, and B is packed into [K/2][2*N] BF16 elements. Each tile still occupies 16 rows of 64 bytes. The packing rule becomes:

1

packed_b[k / 2][2 * n + k % 2] = b[k][n];  // BF16 elements

Enable amx-bf16 for the compute function and store inputs in BF16 encoding. Tile loads copy bytes unchanged, so convert FP32 inputs to BF16 before loading. For validation, promote the same BF16 inputs to FP32, compute a reference result, and compare absolute and relative errors.⁷

Storing: retrieving results from tiles

AMX writes its results to ordinary memory through tile stores. To store the 16×16 INT32 result in tmm0 to a compact two-dimensional array, use:

1
2
3
4
5

alignas(64) std::int32_t c[16][16];
_tile_stored(0, &c[0][0], 16 * sizeof(std::int32_t));
_tile_release();

std::int32_t first = c[0][0];

A tile store also uses the configured row count and byte width, advancing the destination address by stride for each row. The stride above is 64 bytes. If the destination is a block within a larger matrix containing 128 INT32 elements per row, pass 128 * sizeof(std::int32_t) and point base to the block’s top-left element.

After the store, c[0][0] contains the first INT32 element of tmm0’s first row, and c[0][1] contains the next. Scalar or SIMD instructions can then add bias, apply scaling, or evaluate an activation function on the array. _tile_stored() writes matrix data; _tile_storeconfig() writes the 64-byte configuration area.

After storing the results, call _tile_release() to end the current use of tile state. Release itself does not store data. To process more blocks with the same configuration, continue loading, computing, and storing, then release the tiles when the entire computation finishes.

A complete program

Combining these operations gives the program below. Save it as amx_int8.cpp and build and run it with the commands that follow. A uses UINT8, B uses INT8, and both the output and the scalar reference use INT32. The program checks all 256 outputs individually and prints the coordinates and values of any mismatch.

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112

#include <immintrin.h>
#include <cpuid.h>
#include <asm/prctl.h>
#include <sys/syscall.h>
#include <unistd.h>

#include <cstddef>
#include <cstdint>
#include <cstdio>

#ifndef ARCH_REQ_XCOMP_PERM
#define ARCH_REQ_XCOMP_PERM 0x1023
#endif

constexpr int M = 16;
constexpr int N = 16;
constexpr int K = 64;
constexpr unsigned long XFEATURE_XTILEDATA = 18;

struct alignas(64) TileConfig {
    std::uint8_t palette_id;
    std::uint8_t start_row;
    std::uint8_t reserved0[14];
    std::uint16_t colsb[8];
    std::uint8_t reserved1[16];
    std::uint8_t rows[8];
    std::uint8_t reserved2[8];
};
static_assert(sizeof(TileConfig) == 64);
static_assert(offsetof(TileConfig, colsb) == 16);
static_assert(offsetof(TileConfig, rows) == 48);

bool has_amx_int8() {
    unsigned int eax, ebx, ecx, edx;
    if (!__get_cpuid_count(1, 0, &eax, &ebx, &ecx, &edx))
        return false;
    if ((ecx & (1u << 27)) == 0)  // OSXSAVE
        return false;
    if (!__get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx))
        return false;
    const unsigned int required = (1u << 24) | (1u << 25);
    return (edx & required) == required;  // AMX-TILE and AMX-INT8
}

__attribute__((target("amx-tile,amx-int8"), noinline))
void amx_block(const std::uint8_t* a, const std::int8_t* packed_b,
               std::int32_t* c) {
    TileConfig cfg{};
    cfg.palette_id = 1;
    cfg.rows[0] = M;
    cfg.colsb[0] = N * sizeof(std::int32_t);  // C
    cfg.rows[1] = M;
    cfg.colsb[1] = K;                       // A
    cfg.rows[2] = K / 4;
    cfg.colsb[2] = N * 4;                   // packed B

    _tile_loadconfig(&cfg);
    _tile_zero(0);
    _tile_loadd(1, a, K);
    _tile_loadd(2, packed_b, N * 4);
    _tile_dpbusd(0, 1, 2);
    _tile_stored(0, c, N * sizeof(std::int32_t));
    _tile_release();
}

int main() {
    if (!has_amx_int8()) {
        std::fprintf(stderr, "AMX-TILE/INT8 or OSXSAVE unavailable\n");
        return 1;
    }
    if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM,
                XFEATURE_XTILEDATA) != 0) {
        std::perror("ARCH_REQ_XCOMP_PERM");
        return 1;
    }

    alignas(64) std::uint8_t a[M][K];
    alignas(64) std::int8_t b[K][N];
    alignas(64) std::int8_t packed_b[K / 4][N * 4];
    alignas(64) std::int32_t c[M][N]{};
    std::int32_t reference[M][N]{};

    for (int m = 0; m < M; ++m)
        for (int k = 0; k < K; ++k)
            a[m][k] = static_cast<std::uint8_t>((17 * m + 13 * k) % 256);
    for (int k = 0; k < K; ++k)
        for (int n = 0; n < N; ++n)
            b[k][n] = static_cast<std::int8_t>((11 * k + 7 * n) % 31 - 15);

    for (int k = 0; k < K; ++k)
        for (int n = 0; n < N; ++n)
            packed_b[k / 4][4 * n + k % 4] = b[k][n];

    for (int m = 0; m < M; ++m)
        for (int n = 0; n < N; ++n)
            for (int k = 0; k < K; ++k)
                reference[m][n] += static_cast<std::int32_t>(a[m][k]) *
                                   static_cast<std::int32_t>(b[k][n]);

    amx_block(&a[0][0], &packed_b[0][0], &c[0][0]);
    for (int m = 0; m < M; ++m) {
        for (int n = 0; n < N; ++n) {
            if (c[m][n] != reference[m][n]) {
                std::fprintf(stderr, "Mismatch [%d,%d]: %d vs %d\n",
                             m, n, c[m][n], reference[m][n]);
                return 1;
            }
        }
    }
    std::puts("PASS: all 256 outputs match");
    return 0;
}

The compute function follows the seven steps in the data-flow diagram: load the configuration, clear C, load A, load B, compute dot products, store C, and release the tiles. The main function checks CPU features, requests Linux permission, generates inputs, and packs B. It then computes the reference result with three nested loops. Zero-initializing the entire configuration structure leaves reserved fields and unused tiles at zero.

1
2
3

g++ -O3 -std=c++17 -Wall -Wextra amx_int8.cpp -o amx_int8
./amx_int8
objdump -d -Mintel amx_int8 | rg 'ldtilecfg|tilezero|tileloadd|tdpbusd|tilestored|tilerelease'

The function’s target attribute already enables AMX, so this command uses the default global compilation target. A successful check prints PASS: all 256 outputs match. Disassembly lets you verify that the compute function contains the expected tile instructions.

Built with GCC 14.1 on the Xeon Platinum 8458P, this program printed the PASS message above. Extended tests covered UINT8 255, INT8 −128 and 127, blocked accumulation with K=128, zero padding with K=62, and C+=AB. All 104 cases matched an INT64 reference. BF16 packing and dot products also passed 100 tests using inputs that were integer multiples of 0.25, so both the inputs and accumulated results were exactly representable.

If K changes from 64 to 62, the last INT8 group contains only two valid elements. Keep 64 columns in A and the corresponding B packing buffer, fill positions K=62 and 63 with zeros, and reuse the same tile configuration. The padded entries contribute zero products, preserving the result for the first 62 elements. Tile loads read all bytes specified by the configuration, so allocate buffers for the padded dimensions.

Pad K to groups of 4 for INT8 or 2 for BF16. Changes to M or N also require updating the output shape and row stride. A longer K increases the possible integer accumulation range; use an INT64 reference to check whether the INT32 output can overflow.

Measure B packing and AMX computation separately. If B is reused, pack it once and time repeated computations. If every call uses a new B, include packing in the total time. Before timing, warm up the thread that will execute the kernel so that the operating system allocates its AMX state outside the measured interval.⁸

Reusing tiles across output blocks

The amx_block function above loads the configuration and B, performs one tile dot product, and releases the tiles on every call. When computing several output blocks that share B, move the configuration and B load outside the loop. Store A contiguously as [count][16][64] and C as [count][16][16]. Each iteration loads from a new A address and stores to the corresponding C block.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

__attribute__((target("amx-tile,amx-int8"), noinline))
void amx_batch(const std::uint8_t* a, const std::int8_t* packed_b,
               std::int32_t* c, int count) {
    TileConfig cfg{};
    cfg.palette_id = 1;
    for (int t = 0; t < 3; ++t) {
        cfg.rows[t] = 16;
        cfg.colsb[t] = 64;
    }
    _tile_loadconfig(&cfg);
    _tile_loadd(2, packed_b, 64);
    for (int i = 0; i < count; ++i) {
        _tile_zero(0);
        _tile_loadd(1, a + i * 16 * 64, 64);
        _tile_dpbusd(0, 1, 2);
        _tile_stored(0, c + i * 16 * 16, 64);
    }
    _tile_release();
}

This function reuses TileConfig from the earlier example. The same B remains in tmm2 throughout the loop; tmm0 is cleared for each block, and tmm1 receives the next A block. Each A block occupies 1024 bytes, and each C block contains 256 INT32 elements, so the pointers advance by 1024 and 256 elements, respectively. With count=64, the function multiplies 64 blocks of 16×64 A by the same 64×16 B. This is also a matrix multiplication with M=1024, N=16, and K=64. Complete the CPU feature checks and Linux permission request before calling it.

The measurements used GCC 14.1, -O3, and the function target attributes above on a Xeon Platinum 8458P, with one thread pinned to a physical core. A, B, and C were aligned to 64 bytes, and B was packed in advance. Each implementation was warmed up for 2000 calls and measured over 9 rounds with rotating test order. Each round made 1 million calls for the single-block case and 20,000 calls for the other cases. The table reports median time per block, including function calls, configuration, loads, computation, stores, and release. Packing and correctness checks were outside the timed interval, and repeated calls kept the data in cache.

The VNNI baseline performs the same UINT8 × INT8 → INT32 arithmetic on the same packed B. It uses 16 ZMM accumulators for 16 output rows, with each VPDPBUSD instruction computing 16 output columns in one row. The original AMX kernel, the reuse version, and this VNNI kernel all passed elementwise checks on 100 random batches.

For a call that computes just one block, VNNI takes about 57% of the AMX time. Passing 64 blocks that share B to one AMX call makes the reuse version about 3.90× faster than the original and 1.97× faster than this VNNI kernel. Each block performs 16×16×64 multiply-adds. Counting a multiplication and an addition as separate operations, 30.19 ns/block corresponds to about 1.09 TOPS on one core.

The time per block also differs between 7 and 64 blocks. A and C together occupy 14 KiB in the first case and 128 KiB in the second. Increasing the batch size changes both the amortization of configuration overhead and the cache working set. An application can group output blocks that share B into one call, retain B in a tile, and choose the batch size to suit its matrix shape.

A high-performance 32×32 microkernel

In the K=64 example, each output block is stored after a single tile multiply-accumulate. For longer dot products, C can remain in tiles across the entire K loop. A 32×32 output splits into four 16×16 blocks: tmm0 through tmm3 hold the accumulators, tmm4 and tmm5 hold the upper and lower A blocks, and tmm6 and tmm7 hold the left and right B blocks. Each A block contributes to two C blocks, as does each B block.

KTransformers uses this 2×2 output blocking in GemmKernel224Int8 , with configuration and cache blocking along K handled outside the inner computation. Its INT8 kernel accumulates throughout a K block and stores C at the end. The following kernel uses the same division of register roles for UINT8 × INT8 → INT32, with the load sequence written out to show how input reuse works with pipelining.¹⁰

A is packed here as well, so each 16×64 block is contiguous. Its layout is [count][steps][2][16][64], while B uses [steps][2][16][64], with steps=K/64. The count A blocks share one B, and outputs are contiguous in [count][32][32], equivalent to a row-major [count×32][32] matrix. K must be a positive multiple of 64, and input panels are aligned to 64 bytes.

B is reused across adjacent output blocks and uses ordinary tile loads. A advances through the output rows and uses _tile_stream_loadd(). This intrinsic generates TILELOADDT1, which supplies a low-reuse cache hint while retaining the shape and stride semantics of an ordinary load. The hint gives the processor an opportunity to preserve cache space for B; the benefit depends on the working set.⁷

The following function reuses TileConfig. It loads the configuration once before the count loop, clears four C tiles at the start of each output block, and stores them after all K steps are complete.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

__attribute__((target("amx-tile,amx-int8"), noinline))
void amx_gemm32_batch(const std::uint8_t* a, const std::int8_t* b,
                      std::int32_t* c, int steps, int count) {
    TileConfig cfg{};
    cfg.palette_id = 1;
    for (int t = 0; t < 8; ++t) {
        cfg.rows[t] = 16;
        cfg.colsb[t] = 64;
    }
    _tile_loadconfig(&cfg);
    for (int block = 0; block < count; ++block) {
        _tile_zero(0);
        _tile_zero(1);
        _tile_zero(2);
        _tile_zero(3);
        _tile_stream_loadd(4, a, 64);
        _tile_stream_loadd(5, a + 1024, 64);
        _tile_loadd(6, b, 64);
        _tile_loadd(7, b + 1024, 64);
        for (int s = 1; s < steps; ++s) {
            _tile_dpbusd(0, 4, 6);
            _tile_dpbusd(1, 4, 7);
            _tile_stream_loadd(4, a + s * 2048, 64);
            _tile_dpbusd(2, 5, 6);
            _tile_loadd(6, b + s * 2048, 64);
            _tile_dpbusd(3, 5, 7);
            _tile_stream_loadd(5, a + s * 2048 + 1024, 64);
            _tile_loadd(7, b + s * 2048 + 1024, 64);
        }
        _tile_dpbusd(0, 4, 6);
        _tile_dpbusd(1, 4, 7);
        _tile_dpbusd(2, 5, 6);
        _tile_dpbusd(3, 5, 7);
        _tile_stored(0, c, 128);
        _tile_stored(1, c + 16, 128);
        _tile_stored(2, c + 16 * 32, 128);
        _tile_stored(3, c + 16 * 32 + 16, 128);
        a += steps * 2048;
        c += 1024;
    }
    _tile_release();
}

The loop first computes C00 and C01. Both output blocks have now used the current A0, so the next A0 can be loaded into tmm4. Computing C10 completes the second use of the current B0, so the next B0 can be loaded. After computing C11, the kernel replaces A1 and B1. This sequence interleaves upcoming loads with computations on independent accumulators without overwriting inputs still in use. The final four multiply-accumulates outside the loop consume the last loaded panels, avoiding reads beyond K.

For a complete program, reuse the earlier headers, TileConfig, and feature-checking function, and replace main with the code below. It computes M=1024, N=32, and K=1280, packs A and B, and checks all 32768 outputs individually.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

int main() {
    if (!has_amx_int8()) {
        std::fputs("AMX-TILE/INT8 unavailable\n", stderr);
        return 1;
    }
    if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM,
                XFEATURE_XTILEDATA) != 0) {
        std::perror("ARCH_REQ_XCOMP_PERM");
        return 1;
    }

    constexpr int batches = 32, depth = 1280, blocks = depth / 64;
    alignas(64) static std::uint8_t a[batches * 32][depth];
    alignas(64) static std::int8_t b[depth][32];
    alignas(64) static std::uint8_t pa[batches][blocks][2][16][64];
    alignas(64) static std::int8_t pb[blocks][2][16][64];
    alignas(64) static std::int32_t c[batches][32][32];

    for (int m = 0; m < batches * 32; ++m)
        for (int k = 0; k < depth; ++k)
            a[m][k] = (17 * m + 13 * k) % 256;
    for (int k = 0; k < depth; ++k)
        for (int n = 0; n < 32; ++n)
            b[k][n] = (11 * k + 7 * n) % 255 - 128;

    for (int i = 0; i < batches; ++i)
        for (int m = 0; m < 32; ++m)
            for (int k = 0; k < depth; ++k)
                pa[i][k / 64][m / 16][m % 16][k % 64] = a[i * 32 + m][k];
    for (int k = 0; k < depth; ++k)
        for (int n = 0; n < 32; ++n)
            pb[k / 64][n / 16][(k % 64) / 4][4 * (n % 16) + k % 4] = b[k][n];

    amx_gemm32_batch(&pa[0][0][0][0][0], &pb[0][0][0][0],
                     &c[0][0][0], blocks, batches);
    for (int i = 0; i < batches; ++i)
        for (int m = 0; m < 32; ++m)
            for (int n = 0; n < 32; ++n) {
                std::int64_t expected = 0;
                for (int k = 0; k < depth; ++k)
                    expected += std::int64_t(a[i * 32 + m][k]) * b[k][n];
                if (c[i][m][n] != expected) {
                    std::fprintf(stderr, "Mismatch at (%d, %d)\n", i * 32 + m, n);
                    return 1;
                }
            }
    std::puts("PASS: all 32768 outputs match");
}

Save it as amx_gemm32.cpp and compile with the same function-level target attributes:

1
2

g++ -O3 -std=c++17 -Wall -Wextra amx_gemm32.cpp -o amx_gemm32
./amx_gemm32

The complete program passed all output checks on the Xeon Platinum 8458P. A separate test covered count=1, 3, and 32; K=64, 128, and 1280; and inputs containing unsigned 255, signed −128 and 127, or random values. Across these 27 combinations of dimensions and input patterns, both AMX and VNNI matched an INT64 reference element by element, and guard values at both ends of the output buffer remained intact.

The performance tests used GCC 14.1 and one physical core, with its frequency set to 2.2 GHz through Slurm. The ratio of measured core cycles to wall-clock time was also close to 2.2 cycles/ns, allowing AMX and VNNI throughput to be compared at the same frequency. Both used the same packed panels with M=1024 and N=32. Timing included configuration, loads, computation, stores, and release, with packing and validation outside the interval. Each measurement followed 500 warmup calls. The table gives median times for the entire matrix multiplication over 9 rounds with rotating implementation order.

For these measurements, buffers were aligned to 2 MiB, Linux transparent huge pages were requested, and their use was confirmed after touching the memory. Repeatedly used panels were cached. Tests with ordinary pages showed substantial timing variation at K=1280, so record the page size when reproducing the results. Request huge pages after allocation with the following call, where bytes is the allocation size rounded up to a multiple of 2 MiB:

1
2
3
4

#include <sys/mman.h>

// buffer is 2 MiB-aligned; bytes covers the whole allocated region.
madvise(buffer, bytes, MADV_HUGEPAGE);

The VNNI baseline uses 8×32 blocks and 16 ZMM accumulators. Each broadcast of four A bytes updates both halves of an output row. At K=1280, it reaches about 249 integer operations per cycle, close to the processor’s VNNI peak of 256 operations per core per cycle. AMX has an INT8 peak of 2048 operations per core per cycle, giving a same-frequency compute-peak ratio of 8. A multiplication and an addition each count as one operation.¹¹

At K=1280, one call performs 2×1024×32×1280=83,886,080 operations. Its 20.61 µs runtime corresponds to about 4.07 TOPS. The core-cycle count gives about 1851 operations per cycle, or 90.4% of the AMX compute peak, and a 7.43× speedup over the VNNI kernel that is already close to its own peak. At 2.2 GHz, the compute peak implies a compute-only lower bound of about 18.62 µs. The measured 20.61 µs includes configuration, data movement, and output stores.

In the same experiment, ordinary loads for both A and B took 875.35 ns per 32×32 output block. Streaming loads for A reduced this to 648.14 ns; interleaving loads and multiply-accumulates reduced it further to 644.09 ns. Most of the gain came from changing how the inputs use the cache, with instruction scheduling adding about 0.6%.

The B panel occupies 40 KiB at K=1280 and 48 KiB at K=1536. The latter equals this processor’s entire per-core L1 data-cache capacity, which must also serve A, C, and other accesses. Throughput fell as K increased to this size. Blocking must account for both accumulation length and cache capacity: keep C in tiles to reduce intermediate loads and stores, keep repeatedly used B in cache, and advance through A with streaming loads and interleaved scheduling.

Comparison with oneDNN

oneDNN provides a matrix multiplication interface for applications. This comparison used the cluster’s installed oneDNN 3.2.0 native CPU build, with one OpenMP thread, alternating between the library and the microkernel in the same process. The processor, 2.2 GHz frequency, M=1024, N=32, UINT8 × INT8 → INT32 arithmetic, and transparent huge-page conditions matched the preceding tests. For all four K values, every one of the 32,768 outputs from both implementations matched an INT64 reference. The oneDNN log reported brg:avx512_core_amx and one thread.

oneDNN receives row-major A and selects B’s internal layout through any in the weights descriptor. Primitive creation, scratchpad allocation, and a one-time weight reorder occur before timing; the measured interval covers execute and stream.wait. The microkernel uses prepacked A and B, and its entire function call is timed. Both implementations repeatedly use their own input buffers. Each measurement follows 500 warmup calls, and the reported value is the median of 9 rounds with alternating test order. The oneDNN layout selection looks like this:¹²

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

using dt = dnnl::memory::data_type;
using tag = dnnl::memory::format_tag;

dnnl::primitive_attr attr;
attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
auto pd = dnnl::matmul::primitive_desc(
    engine,
    dnnl::memory::desc({M, K}, dt::u8, tag::ab),
    dnnl::memory::desc({K, N}, dt::s8, tag::any),
    dnnl::memory::desc({M, N}, dt::s32, tag::ab),
    attr);
// Reorder B into pd.weights_desc() once before repeated execution.

These results compare call latency for fixed N=32 with reused inputs. At K=1280, this kernel’s B panel occupies 40 KiB, while oneDNN’s chosen weight layout occupies 80 KiB. The cache working sets also differ with the layouts. The specialized microkernel fixes the output width, panel layout, and register allocation together, reducing runtime for this shape. If A changes on every request, include its packing time in request latency. If matrix dimensions, data types, or fused operations change frequently, oneDNN’s layout selection and API can save the corresponding kernel development work.

Summary

The basic example uses three tiles to load inputs, compute dot products, and store results. The 32×32 microkernel uses four accumulators to retain intermediate results while reusing two A blocks and two B blocks. The fields rows, colsb, and stride determine the load shape; B’s packing determines which elements participate in each multiply-accumulate. Performance also depends on accumulation length, input reuse in cache, and load/compute scheduling. In these measurements, VNNI suits an isolated small block, while larger matrix blocks sharing B benefit from AMX throughput. When choosing a block size, examine operations per cycle, the cache working set, and full-call latency together to identify the bottleneck.