C/C++ 中使用 Intel AMX：寄存器、指令接口与数据流

Intel AMX（Advanced Matrix Extensions）是 x86 处理器面向矩阵运算的一组指令扩展。它增加了二维 tile 寄存器，可以把两个矩阵块的乘积累加到第三个矩阵块中。神经网络的全连接层、转为矩阵乘法的卷积，以及批量向量点积，都可以用这种形式组织内层计算。¹

AMX 从第 4 代 Xeon Scalable（Sapphire Rapids）开始提供，首批支持 INT8 和 BF16；第 5 代 Xeon Scalable（Emerald Rapids）延续了这两种能力。Xeon 6 的 P-core 产品（Granite Rapids）也支持 AMX，并增加了 FP16；E-core 产品（Sierra Forest）不支持 AMX。工作站产品中也有支持 AMX 的型号，例如 Xeon W-3400 系列的 w9-3495X。本文使用的 Xeon Platinum 8458P 属于 Sapphire Rapids，下面的 INT8 示例适用于具备 AMX-TILE 和 AMX-INT8 的处理器。²³⁴

C/C++ 程序通过 intrinsics 使用这些指令。计算内存中两个矩阵的乘积 C=AB 时，程序把 A 和 B 的子块读入 tile，完成乘加后将 C 写回数组。加载和存储访问的都是进程中的普通内存，也经过 CPU 的缓存。

下面以 UINT8 × INT8 → INT32 为例，固定计算一个 16×16 的输出块，并在 Xeon Platinum 8458P 上验证结果与性能。每段代码都对应 A、B、C 的一次加载、计算或写回，说明数据在内存和 tile 中的位置如何变化。

头文件与编译选项

AMX intrinsics 通过 immintrin.h 引入，和 AVX、AVX-512 使用同一个入口头文件：

1
#include <immintrin.h>

一个 INT8 内核同时用到两类指令：AMX-TILE 配置和搬运 tile，AMX-INT8 执行整数点积。编译器需要同时启用这两个扩展。使用支持 AMX 的 GCC 工具链时，命令如下：⁵

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

使用 BF16 输入、FP32 累加时，还需要 -mamx-bf16。编译标志决定允许生成哪些指令；运行程序的 CPU 必须提供这些扩展，Linux 也必须允许进程使用 tile 状态。

还可以只为计算函数启用 AMX。下面的 target 属性为计算函数启用 AMX，main 和能力检查仍按普通 x86-64 目标编译。文末的完整示例也采用这种写法。

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 和 Clang 都提供下文使用的显式 tile 编号接口，定义可在 GCC 与 LLVM 的头文件中查阅。本文的编译命令使用 GCC。遇到未知 intrinsic、target option 或汇编指令时，应同时检查编译器与汇编器版本。⁶⁷

在 Linux 上，还需要在第一次使用 tile 数据前向内核申请权限：

 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;
}

这里的 18 是扩展状态编号。内核据此允许进程使用 tile 数据，并在执行线程首次使用时分配保存状态所需的空间。权限在进程内有效，fork 会继承，exec 会清除。如果应用安装了备用信号栈，还要让它容纳增大的信号帧。完整示例会在申请权限前检查 CPU 的 AMX-TILE、AMX-INT8 和 OSXSAVE 特性。⁸

寄存器与核心接口

AMX 使用 TILECFG 保存各个 tile 的形状，使用 tmm0 到 tmm7 保存参与计算的数据。本文采用的 palette 1 提供 8 个 tile，每个最多 16 行、每行 64 字节，合计最多 8 KiB 数据。¹

图 1：AMX 的配置状态与 tile 寄存器
TILECFG 配置形状，8 个 tile 保存数据。每行最多 64 字节，换成不同元素类型后，列数也会改变。缩略图每格代表 4 字节，放大图细格代表 1 字节。本文绘制。

配置字段 rows 指定行数，colsb 指定每行的字节数。将两者设为 16 和 64，就得到一个 1024 字节的 tile。用于 INT8 时，每行有 64 个元素；用于 BF16 时，每行有 32 个元素；用于 INT32 时，每行有 16 个元素。三种情况下配置可以完全相同，计算指令决定如何解释这些字节。

下面列出配置、计算和搬运数据所需的接口。dst、a、b 是 tile 编号，base 是内存地址，stride 是相邻内存行起点之间的字节距离。

接口	作用	参数含义
`_tile_loadconfig(&cfg)`	加载 tile 配置	cfg 指向 64 字节配置区
`_tile_storeconfig(&cfg)`	将当前配置写回内存	用于读取配置状态
`_tile_zero(dst)`	清零目标 tile	常用于初始化累加器
`_tile_loadd(dst, base, stride)`	从内存逐行加载 tile	形状来自配置，行跨度来自 stride
`_tile_dpbusd(dst, a, b)`	UINT8 × INT8，累加到 INT32	dst 同时是输入累加器与输出
`_tile_dpbssd(dst, a, b)`	INT8 × INT8，累加到 INT32	两个输入均按有符号整数解释
`_tile_dpbf16ps(dst, a, b)`	BF16 × BF16，累加到 FP32	K 方向每两个 BF16 为一组
`_tile_stored(src, base, stride)`	将 tile 逐行写回内存	stride 的单位仍是字节
`_tile_release()`	结束 tile 使用并恢复初始状态	后续计算需要重新加载配置

这些接口直接指定寄存器编号。例如，_tile_zero(0) 清零 tmm0，_tile_zero(1) 清零 tmm1，编号必须在编译期确定。因此，内核通常固定使用几个 tile，循环中改变加载地址，让同一组寄存器依次处理不同的矩阵块。⁶

计算 $C_{M\times N} += A_{M\times K}B_{K\times N}$ 时，固定使用 tmm0 保存 C，tmm1 保存 A，tmm2 保存打包后的 B。INT8 示例取 M=16、N=16、K=64，三个 tile 都是 16 行、每行 64 字节，但它们的元素类型与数学含义不同。

图 2：一次 INT8 矩阵乘法的数据流
A 与打包后的 B 从内存进入 tmm1、tmm2，点积累加到 tmm0，最后写回 INT32 数组。本文绘制。

读入：配置形状与内存布局

加载数据前，先配置 tile 的形状。下面的结构体对应配置区布局：palette_id 位于字节 0，colsb 从偏移 16 开始，rows 从偏移 48 开始。未使用的字段保持为 0。¹

 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);

现在把 A[16][64] 放入 tmm1。A 的元素为 UINT8，每行占 64 字节，下一行紧接着上一行。配置已经指定读取 16 行、每行 64 字节，调用中只需给出首地址和 64 字节的行跨度：

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

从大矩阵读取子块时，colsb 和 stride 通常不同。例如，A 每行有 128 个 UINT8，要从第 m0 行、第 k0 列开始读取一个 16×64 子块，colsb 设为 64，stride 设为 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);

这次加载读取每行的 64 字节，再跳到相隔 128 字节的下一行起点。每行剩余的 64 字节不进入 tile。读完 16 行后，tmm1 中的数据仍是紧凑的 16×64 字节，内存中的行间空隙不会保存在寄存器中。

接着读取 B。计算 C[m][n] 需要 A 的第 m 行与 B 的第 n 列做点积，但行主序 B 把同一行的不同列放在一起。AMX 的 INT8 指令按 4 对字节做乘加，需要将 B 同一列的 4 个连续 K 元素排在一起。下面的循环完成这个重排：

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];

图 3：B 沿 K 方向每 4 个元素打包
同一颜色对应同一输出列。每行先放一列的 4 个 K 元素，再放下一列的 4 个 K 元素；下标保留原矩阵坐标。本文绘制。

图中以 K=8、N=4 为例。打包结果的第一行依次放入第 0 列的 b00、b10、b20、b30，再放入第 1 列的 b01、b11、b21、b31，随后依次放入第 2、3 列各自的四个元素。下一行处理 K 方向上的后 4 个元素。颜色区分输出列，元素下标则说明每个字节来自原矩阵的哪个位置。

回到 K=64、N=16，打包后的 B 有 K/4=16 行，每行 4N=64 字节，可以直接读入 tmm2：

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

tile load 按行复制字节，因此 B 的重排在加载前完成。如果直接加载原始 B，寄存器里相邻的字节仍来自不同的输出列；计算指令却会把它们当成同一列的 4 个 K 元素相乘。这种错误的元素配对可能不会触发运行异常，却会改变计算结果。⁷

计算：点积与累加器

数据就位后，计算只需要指定三个 tile：

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

第一条语句令 tmm0 中的累加值为 0。第二条读取 tmm1 和 tmm2，将点积结果加到 tmm0 中；两个输入 tile 的内容不变。对于输出位置 (m,n)，参与运算的是 A 的第 m 行，以及打包后 B 中属于第 n 列的各组字节。这个操作可以展开成下面的标量循环：

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]);

以一组数为例：A 的前 4 个元素为 [1, 2, 3, 4]，打包后的 B 中，输出第 0 列对应 [−1, 2, −3, 4]，第 1 列对应 [−2, 1, 0, 3]。这组数据对两个输出的累加值分别为：

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

后续 K 组继续累加到这两个位置。一次 tile 点积会对配置覆盖的所有行和输出列执行相同的操作。

这里每组 4 对 INT8 相乘，累加到对应的 INT32 元素。A 为 UINT8、B 为 INT8 时使用 dpbusd；两个输入都为 INT8 时使用 dpbssd。内存里的 0xff 在两种解释下分别是 255 和 −1，所以测试输入应包含大于 127 的无符号值和负的有符号值。⁷

tmm0 的旧值会参与累加。要计算 C=AB，可以清零后执行一次或多次点积；要计算 C+=AB，就先把原有的 C 读入 tmm0：

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

当 K 超过单个 tile 能容纳的范围时，可以让 C 一直留在 tmm0，依次加载 A 和 B 的下一组 K 块。以下片段假设 A 有 128 列，B 已打包成两个独立的 64×16 面板，每个面板的存储形状为 [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.

配置加载放在累加循环外，循环内只加载输入和执行点积。重新加载配置会初始化 tile 数据状态，因此累加期间保持同一份配置。⁹如果需要多个输出块，可以给每块分配一个 C tile，复用已加载的 A 或 B。输入块和输出块共同占用这 8 个 tile。

BF16 的数据流相同，但一组点积包含两个 BF16，结果累加到 FP32。对应接口是：

1
_tile_dpbf16ps(0, 1, 2);

若选择 M=16、N=16、K=32，A 占 16×32 个 BF16，C 占 16×16 个 FP32，B 打包为 [K/2][2*N] 个 BF16。三个 tile 仍各占 16 行、每行 64 字节。B 的打包关系改为：

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

计算函数需要启用 amx-bf16，输入缓冲区按 BF16 编码存放数据。tile load 原样搬入字节；FP32 输入应在加载前转换为 BF16。验证时应先将同一份 BF16 输入提升到 FP32，计算参考结果，再比较绝对误差与相对误差。⁷

写回：从 tile 取出结果

AMX 的结果通过 tile store 写回普通内存。对 tmm0 中的 16×16 INT32 结果，写回紧凑二维数组可以这样做：

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];

tile store 同样使用配置里的行数和行字节数，并按 stride 移动目标地址。上面的 stride 是 64 字节。如果输出落在每行 128 个 INT32 的大矩阵中，写回子块时应传入 128 * sizeof(std::int32_t)，并把 base 移到目标子块的左上角。

执行 store 后，c[0][0] 就是 tmm0 第一行的第一个 INT32 元素，c[0][1] 是下一个元素。要给输出加偏置、做缩放或激活，可以在这个数组上使用标量或 SIMD 指令。_tile_stored() 写回矩阵数据，_tile_storeconfig() 则写回 64 字节配置区。

结果写回后，调用 _tile_release() 结束当前 tile 使用。release 本身不执行存储。如果还要用相同配置处理更多块，可以继续加载、计算和写回，在整段计算结束时再释放 tile。

把数据流连成一个程序

将前面的操作组合起来，就得到下面的程序。保存为 amx_int8.cpp，用下方命令编译运行。A 的类型为 UINT8，B 的类型为 INT8，输出和标量参考都使用 INT32。程序逐一比较全部 256 个输出，出现不一致时打印具体坐标和值。

  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;
}

计算函数里的七步就是前面图示的数据流：加载配置、清零 C、加载 A、加载 B、点积、写回 C、释放 tile。main 负责能力检查、Linux 权限申请、数据生成和 B 打包，随后用普通三重循环算出参考结果。配置结构体整体清零，保证保留字段和未使用的 tile 保持为 0。

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'

函数的 target 属性已启用 AMX，这条命令使用默认的全局编译目标。程序校验成功时输出 PASS: all 256 outputs match。反汇编则可以检查计算函数是否生成了预期的 tile 指令。

在 Xeon Platinum 8458P 上，GCC 14.1 编译的这个程序输出了上述 PASS。扩展测试还检查了 UINT8 的 255、INT8 的 −128 和 127、K=128 的分块累加、K=62 的补零，以及 C+=AB，共 104 组，均与 INT64 参考结果一致。BF16 的打包和点积通过了 100 组测试，输入取 0.25 的整数倍，使输入和累加结果都能精确表示。

如果把 K 从 64 改成 62，最后一组 INT8 点积只剩两个有效元素。可以仍分配 64 列的 A 和相应的 B 打包区，将 K=62、63 的位置补零，继续使用同一个 tile 配置。补入的元素产生零乘积，保留前 62 个元素的计算结果。tile load 会读取配置指定的全部字节，因此缓冲区要按补齐后的大小分配。

INT8 按 4 个 K 元素分组，BF16 按 2 个分组，补齐时据此确定大小。改变 M、N 时，也需要同步调整输出形状和行跨度。延长 K 还会增加整数累加值的范围，可以用 INT64 计算参考结果，检查 INT32 输出是否会溢出。

测量性能时，将 B 打包和 AMX 计算分别计时。如果 B 会被重复使用，可以预先打包，再测量多次计算；如果每次调用都换 B，打包成本也应计入总时间。开始测量前让实际执行内核的线程完成一次预热，使首次使用 AMX 的状态分配发生在计时之外。⁸

连续计算时复用 tile

上面的 amx_block 每次调用都加载配置、读入 B、执行一次点积，然后释放 tile。连续计算多个共享 B 的输出块时，可以把配置和 B 的加载移到循环外。A 按 [count][16][64] 连续存放，C 按 [count][16][16] 存放；每次迭代只更换 A 的地址，并把结果写到对应的 C 块。

 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();
}

这个函数沿用前面的 TileConfig。tmm2 在整个循环中保存同一个 B，tmm0 每次清零，tmm1 加载下一个 A 块。每块 A 占 1024 字节，每块 C 占 256 个 INT32，因此两个指针分别前进 1024 和 256 个元素。传入 count=64 时，计算的是 64 个 16×64 的 A 块与同一个 64×16 的 B 相乘，也可以看成 M=1024、N=16、K=64 的矩阵乘法。调用前仍需完成 CPU 检查和 Linux 权限申请。

在 Xeon Platinum 8458P 上，用 GCC 14.1、-O3 和上述函数级 target 属性编译，固定单线程到一个物理核。A、B、C 缓冲区按 64 字节对齐，B 预先打包，计时前预热 2000 次。每种实现测量 9 轮，轮换测试顺序；单块场景每轮调用 100 万次，其余场景每轮调用 2 万次。下表给出每块耗时的中位数，包含函数调用、配置、加载、计算、写回和释放。B 打包与结果校验放在计时之外，数据在重复调用中保持热缓存。

VNNI 对照使用相同的 UINT8 × INT8 → INT32 运算和同一份打包后的 B，以 16 个 ZMM 累加器分别计算 16 行输出，每条 VPDPBUSD 计算一行的 16 个输出列。原版、复用版和这个 VNNI 内核都通过了 100 组随机批次的逐元素校验。

每次调用的块数	原版 AMX（ns/块）	复用配置与 B 的 AMX（ns/块）	VNNI 对照（ns/块）
1	104.21	104.69	59.33
3	104.13	38.66	57.80
4	104.04	30.28	57.65
7	103.98	23.43	58.13
64	117.65	30.19	59.45

一次调用只计算一个块时，VNNI 的耗时约为 AMX 的 57%。把 64 个共享 B 的块交给一次 AMX 调用后，复用版本相对原版加速约 3.90 倍，相对上述 VNNI 内核加速约 1.97 倍。每块包含 16×16×64 次乘加，按一次乘法和一次加法各计一个操作，30.19 ns/块对应约 1.09 TOPS 的单核吞吐。

7 个块与 64 个块的每块耗时也不同。前者的 A 和 C 合计 14 KiB，后者为 128 KiB；增大批次时，配置开销的摊销和缓存工作集会同时变化。实际程序可以把共享 B 的输出块放在同一次调用里，让 B 留在 tile 中，同时根据矩阵形状选择批次大小。

32×32 高性能微内核

前面的 K=64 示例，每个输出块只做一次矩阵乘加就写回。计算更长的点积时，可以让 C 跨整个 K 循环留在 tile 中。一次计算 32×32 的输出，正好分成四个 16×16 子块：tmm0 到 tmm3 保存这四个累加器，tmm4、tmm5 保存上下两块 A，tmm6、tmm7 保存左右两块 B。每块 A 用于两个 C，每块 B 也用于两个 C。

KTransformers 的 GemmKernel224Int8 采用了这种 2×2 输出分块，并将配置与 K 方向的缓存分块放在外层。它的 INT8 内核在一个 K 分块内持续累加，结束时才写回 C。下面用同样的寄存器分工，继续计算 UINT8 × INT8 → INT32，并把加载顺序展开，观察输入复用和流水线如何配合。¹⁰

四个 C tile 跨 K 循环累加，两块 A 与两块 B 交叉组合。 — 图 4：32×32 输出块的八个 tile 分工
四个 C tile 跨 K 循环累加。图中 B 标注数学形状，装入寄存器前仍按每四个 K 元素打包成 16×64 字节。本文绘制。

这里也将 A 打包，使每个 16×64 子块连续存放。A 的布局为 [count][steps][2][16][64]，B 的布局为 [steps][2][16][64]，其中 steps=K/64。count 个 A 块共享同一个 B，输出连续存放为 [count][32][32]，也就是一个行主序的 [count×32][32] 矩阵。K 为正的 64 的倍数，输入面板按 64 字节对齐。

B 会在相邻输出块之间反复使用，采用普通 tile load；A 随输出行向前推进，使用 _tile_stream_loadd()。后者生成 TILELOADDT1，为内存访问提供低复用的缓存提示，加载形状和 stride 的含义与普通 load 相同。这个提示让处理器有机会为反复使用的 B 保留缓存空间；具体收益由工作集决定。⁷

下面的函数沿用前文的 TileConfig。配置在整个 count 循环外加载一次，四个 C tile 在每个输出块开始时清零，完成所有 K 步长后写回。

 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();
}

循环先计算 C00 和 C01，此时旧的 A0 已被两个输出消费，可以把下一块 A0 装入 tmm4。接着计算 C10，旧的 B0 也完成了两次使用，便可加载下一块 B0。最后计算 C11，再替换 A1 和 B1。这样将后续加载穿插在独立累加器的计算之间，并避免覆盖仍要使用的输入。循环外的最后四条乘加处理已加载的末尾面板，因此不会读取 K 范围之外的数据。

完整调用可以复用前文的头文件、TileConfig 和能力检查函数，将 main 替换为下面的代码。它计算 M=1024、N=32、K=1280，完成 A、B 打包后，逐一检查全部 32768 个输出。

 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");
}

保存为 amx_gemm32.cpp，使用同样的函数级 target 属性编译：

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

在 Xeon Platinum 8458P 上，上面的完整程序通过了全部输出校验。另一个测试覆盖 count=1、3、32，K=64、128、1280，以及无符号 255、有符号 −128、127 和随机输入，共 27 组尺寸与数据组合；AMX 和 VNNI 都用 INT64 参考结果逐元素检查，并检查输出缓冲区两端的哨兵值。

性能测试使用 GCC 14.1，固定一个物理核，并通过 Slurm 将频率设为 2.2 GHz。核心周期计数与墙钟时间的比值也接近 2.2 周期/ns。这样可以直接比较 AMX 与 VNNI 在相同频率下的吞吐。两者使用同一份打包面板，M=1024、N=32，计时包含配置、加载、计算、写回与释放；打包和校验在计时之外。每轮先预热 500 次，测量 9 轮并轮换实现顺序，表中为整次矩阵乘法耗时的中位数。

这一组计时将缓冲区按 2 MiB 对齐，申请 Linux 透明大页，并在访问后确认大页实际生效。反复使用的面板位于缓存中。普通页下，K=1280 的重复测试出现过明显的耗时波动，复现时应记录页大小。大页可在分配后通过下面的调用申请，bytes 是向上取整到 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);

VNNI 对照采用 8×32 子块，以 16 个 ZMM 寄存器保存累加值，一次广播 A 的四个字节，同时更新左右两组输出。K=1280 时，它达到约 249 次整数操作/周期，接近该处理器每核 256 次操作/周期的 VNNI 峰值。AMX 的 INT8 峰值为每核 2048 次操作/周期，两者的同频计算峰值之比为 8。这里均将一次乘法和一次加法分别计为一个操作。¹¹

K	VNNI 整次耗时（µs）	AMX 整次耗时（µs）	AMX 加速比	AMX 计算峰值利用率
512	61.77	9.34	6.61×	79.8%
1024	122.92	16.90	7.27×	88.2%
1280	153.04	20.61	7.43×	90.4%
1536	183.48	30.43	6.03×	73.5%

K=1280 时，一次调用完成 2×1024×32×1280=83,886,080 次操作，20.61 µs 对应约 4.07 TOPS。核心周期计数给出约 1851 次操作/周期，即 AMX 计算峰值的 90.4%；相对已接近自身峰值的 VNNI 内核，加速比为 7.43 倍。在 2.2 GHz 下，按计算峰值推得的纯计算时间下界约为 18.62 µs，实测为 20.61 µs，包含配置、数据搬运和输出写回。

在同一轮对照中，A、B 都使用普通加载时，每个 32×32 输出块耗时 875.35 ns。将 A 改为流式加载后降到 648.14 ns，再交错安排加载与乘加后为 644.09 ns。这次收益主要来自输入的缓存使用方式，指令顺序调整带来了约 0.6% 的额外改善。

K=1280 的 B 面板占 40 KiB；K=1536 时增至 48 KiB，已经等于这颗处理器的单核 L1 数据缓存容量，A、C 及其他访问也需要缓存空间。实测中，继续增大 K 使吞吐回落。分块要同时照顾寄存器累加长度和缓存容量：让 C 留在 tile 中减少中间结果存取，让反复使用的 B 留在缓存中，再用流式加载和交错调度推进 A。

与 oneDNN 对照

oneDNN 提供了可以直接用于应用的矩阵乘法接口。这里使用集群安装的 oneDNN 3.2.0 原生 CPU 版本，设置 OpenMP 单线程，在同一进程中交替测量它与上面的微内核。处理器、2.2 GHz 频率、M=1024、N=32、UINT8 × INT8 → INT32 和透明大页条件均与前面的测试一致。四种 K 长度下，两种实现的全部 32,768 个输出都与 INT64 参考计算一致。oneDNN 的运行日志显示实现为 brg:avx512_core_amx，线程数为 1。

oneDNN 接收行主序 A，并通过 weights 描述符中的 any 选择 B 的内部布局。创建 primitive、分配 scratchpad 和一次性权重 reorder 都在计时前完成，计时覆盖 execute 与 stream.wait。本文微内核使用预先打包的 A、B，计时覆盖整个函数调用。两者重复使用各自的输入缓冲区，每轮预热 500 次，轮换测量顺序，取 9 轮的中位数。oneDNN 的布局选择方式如下：¹²

 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.

K	oneDNN AMX（µs）	本文微内核（µs）	微内核相对 oneDNN 的加速比
512	16.52	9.37	1.76×
1024	27.29	16.92	1.61×
1280	32.45	20.63	1.57×
1536	38.31	30.47	1.26×

这组结果比较的是固定 N=32、输入重复使用时的调用耗时。K=1280 时，本文的 B 面板占 40 KiB，oneDNN 选择的权重布局占 80 KiB；两者的缓存工作集也随布局发生变化。专用微内核将输出宽度、面板布局和寄存器分配一起固定下来，因而能在这个形状上进一步压低耗时。应用中若 A 每次更新，还需要把 A 的打包时间计入请求；若矩阵尺寸、数据类型或融合操作经常变化，oneDNN 的布局选择与通用接口可以省去相应的内核开发工作。

小结

基本示例用三个 tile 完成加载、点积和写回，32×32 微内核则用四个累加器保留中间结果，并复用两块 A 与两块 B。加载形状由 rows、colsb 和 stride 决定，乘加关系由 B 的打包布局决定；性能还取决于累加长度、输入在缓存中的复用，以及加载与计算的调度。实测中，单个小块适合 VNNI，共享 B 的较大矩阵块能发挥 AMX 的吞吐优势。选择分块时，把每周期吞吐、缓存工作集和完整调用耗时一起检查，才能看清瓶颈。

头文件与编译选项

寄存器与核心接口

图 1：AMX 的配置状态与 tile 寄存器

图 2：一次 INT8 矩阵乘法的数据流