1. 从一个朴素问题开始：CPU 为什么需要 AVX/AMX？#

我们平时写程序时，脑子里常见的模型是：

1
for (int i = 0; i < n; ++i) {
2
    c[i] = a[i] + b[i];
3
}

这段代码看起来像是 CPU 一次处理一个 a[i] + b[i]。如果真是这样，CPU 每执行一次循环，只得到一个结果。

但现代 CPU 并不满足于”一次只算一个数”。很多任务天然有大量重复计算，例如：

• 图片处理：一批像素做同样的加减乘除。
• 音视频编解码：一批采样点做滤波、变换。
• 科学计算：一批浮点数做向量和矩阵运算。
• 机器学习：矩阵乘法、卷积、attention 里有大量乘加。

于是 CPU 里出现了面向数据并行的指令和执行单元。

SIMD 是 Single Instruction, Multiple Data，意思是：一条指令，同时处理多个数据。

AVX 和 AMX 都属于提高数据并行吞吐量的硬件扩展，只是抽象层次不同：

• AVX：Advanced Vector Extensions，是典型的 SIMD，主要面向一维向量计算。
• AMX：Advanced Matrix Extensions，主要面向二维矩阵块计算，尤其服务于 AI 里的矩阵乘法。

它们不是新的编程语言，也不是某个库，而是 CPU 指令集架构的一部分。软件要用它们，必须满足三件事：

1. CPU 硬件真的支持这些指令。
2. 操作系统愿意保存和恢复这些额外的寄存器状态。
3. 编译器、运行库或程序员生成了对应指令。

这三层关系可以画成下面这样：

2. 指令集、寄存器和执行单元#

要理解 AVX/AMX，先把几个概念分开。

指令集（ISA） 是 CPU 对软件承诺的”可用操作清单”。比如 add 是整数加法，vmulps 是 AVX 浮点向量乘法，tdpbf16ps 是 AMX BF16 tile 矩阵乘加。程序最终都会变成这些机器指令。

寄存器 是 CPU 内部最快的一小块存储。普通整数代码用通用寄存器，AVX 用 XMM/YMM/ZMM 这类向量寄存器，AMX 用 tile 寄存器。

执行单元 是真正干活的电路。指令被解码后，CPU 会把它送到合适的执行单元，例如整数 ALU、浮点单元、向量 FMA 单元，或者 AMX 的 TMUL 矩阵乘法单元。

所以，“支持 AVX”并不只是多了几个指令名字。它通常意味着 CPU 还提供了更宽的向量寄存器、更复杂的执行单元，以及相应的上下文保存机制。

3. AVX 是什么？#

AVX 是 x86 平台上的向量指令扩展。它的核心思想是：把多个相同类型的数据打包进一个向量寄存器，然后用一条指令统一处理。

举例说，如果我们有 8 个 float：

1
a = [a0, a1, a2, a3, a4, a5, a6, a7]
2
b = [b0, b1, b2, b3, b4, b5, b6, b7]

普通标量代码会做 8 次加法。AVX 可以把这 8 个 float 放进一个 256-bit 的 YMM 寄存器，用一条向量加法得到 8 个结果。

常见 AVX 家族大致可以这样理解：

为了让这些名字更具体一点，可以看几条常见指令。下面的表不追求完整，只是帮助建立感觉：

从命名上也能看出一些规律：PS 常表示 packed single-precision float，PD 表示 packed double，D 常和 32-bit doubleword 整数有关。前缀 V 通常说明这是 VEX/EVEX 编码下的向量指令。

需要注意：AVX 不是一个单一能力开关。现实中经常会看到 AVX、AVX2、FMA、AVX512F、AVX512VNNI 等多个 feature flag。软件通常要检查自己需要的那一组，而不是只问”有没有 AVX”。

4. AVX 和软件的关系#

软件使用 AVX 通常有三条路。

1
void add_arrays(const float* a, const float* b, float* c, int n) {
2
    for (int i = 0; i < n; ++i) {
3
        c[i] = a[i] + b[i];
4
    }
5
}

1
g++ -O3 -mavx2 add.cpp -o add

1
#include <immintrin.h>
2

3
void add8_avx(const float* a, const float* b, float* c) {
4
    __m256 va = _mm256_loadu_ps(a);
5
    __m256 vb = _mm256_loadu_ps(b);
6
    __m256 vc = _mm256_add_ps(va, vb);
7
    _mm256_storeu_ps(c, vc);
8
}

1
import numpy as np
2

3
a = np.random.rand(1_000_000).astype(np.float32)
4
b = np.random.rand(1_000_000).astype(np.float32)
5
c = a + b

5. AVX 的限制：不是越宽越无脑快#

AVX 很有用，但它不是魔法。

第一，SIMD 喜欢规整的数据。连续数组、相同操作、少分支，是 SIMD 最舒服的场景。如果数据结构是链表、树、哈希表，或者每个元素走不同分支，AVX 很难发挥。

第二，性能可能被内存带宽限制。如果计算很简单，比如只是 a[i] + b[i]，CPU 可能不是卡在加法，而是卡在从内存读写数据。

第三，AVX-512 这类宽向量在某些处理器上可能影响频率。具体影响和微架构、负载、散热、功耗策略有关，不能简单理解成”512-bit 一定比 256-bit 快两倍”。

第四，代码需要处理兼容性。如果你用 -mavx2 编译整个程序，生成的二进制在不支持 AVX2 的 CPU 上可能直接非法指令崩溃。工程上通常会做 runtime dispatch：启动时检测 CPU feature，然后选择对应实现。

6. AMX 是什么？#

AMX 是 Intel 面向矩阵计算引入的扩展。和 AVX 主要处理一维向量不同，AMX 引入了二维的 tile 寄存器。

根据 Intel 对 AMX 的介绍，AMX 的核心由两部分组成：

• Tiles：每个核心上的二维 tile 寄存器。Intel 文档中描述为 8 个 tile，每个 tile 最大 1 KiB。
• TMUL：Tile Matrix Multiplication，负责对 tile 做矩阵乘法相关计算的硬件单元。

直观上可以把 AMX 想成：CPU 核心里多了一组”小矩阵暂存区”，以及专门吃这些小矩阵的乘加单元。

AMX 的目标非常明确：加速深度学习训练和推理中常见的矩阵乘法。它支持的具体数据类型随扩展而定，例如 AMX-BF16、AMX-INT8、AMX-FP16 等。

AMX 的指令比 AVX 更有”操作 tile”的味道。常见指令可以粗略分成三类：

如果说 AVX 的典型指令像是”对 8 个或 16 个数一起做加法/乘法”，AMX 的典型指令更像是”把 A 的一个小矩阵块和 B 的一个小矩阵块相乘，并累加到 C 的一个小矩阵块”。这也是为什么 AMX 更适合 GEMM、卷积、Transformer 里那些最终会落到矩阵乘法的计算。

这和 AVX/AVX-512 的关系不是”谁替代谁”，而是抽象层次不同：

7. AMX 如何使用？#

AMX 比 AVX 更不像”随手写两行就能用”。它至少涉及三层准备。

1
lscpu | grep -i amx

1
grep -o 'amx[^ ]*' /proc/cpuinfo | sort -u

1
#include <immintrin.h>
2
#include <stdint.h>
3

4
struct tile_config {
5
    uint8_t palette_id;
6
    uint8_t start_row;
7
    uint8_t reserved[14];
8
    uint16_t colsb[16];
9
    uint8_t rows[16];
10
};
11

12
void configure_tiles() {
13
    tile_config cfg = {};
14
    cfg.palette_id = 1;
15

16
    // tmm0/tmm1/tmm2 的行数和每行字节数。
17
    // 真实代码要根据矩阵分块、数据类型和指令要求仔细设置。
18
    cfg.rows[0] = 16;
19
    cfg.colsb[0] = 64;
20
    cfg.rows[1] = 16;
21
    cfg.colsb[1] = 64;
22
    cfg.rows[2] = 16;
23
    cfg.colsb[2] = 64;
24

25
    _tile_loadconfig(&cfg);
26
}

1
import torch
2

3
model = model.eval().to("cpu")
4
inputs = inputs.to("cpu")
5

6
with torch.no_grad(), torch.autocast(device_type="cpu", dtype=torch.bfloat16):
7
    outputs = model(inputs)

8. 如何判断自己的机器支持什么？#

在 Linux 上可以先看 CPU flags：

1
lscpu

或者：

1
grep -m1 flags /proc/cpuinfo

你可能会看到类似：

1
avx avx2 fma avx512f avx512_vnni amx_bf16 amx_tile amx_int8

对 C/C++ 程序来说，更可靠的做法是在运行时检测 feature。GCC/Clang 提供了类似接口：

1
#include <iostream>
2

3
int main() {
4
#if defined(__x86_64__) || defined(__i386__)
5
    __builtin_cpu_init();
6

7
    std::cout << "avx2: "
8
              << __builtin_cpu_supports("avx2") << "\n";
9
    std::cout << "avx512f: "
10
              << __builtin_cpu_supports("avx512f") << "\n";
11
#endif
12
}

工程化一点的程序会准备多个实现：

• 标量 fallback：任何机器都能跑。
• SSE/AVX2 实现：覆盖较多桌面和服务器 CPU。
• AVX-512 实现：用于支持更宽向量的 CPU。
• AMX 实现：用于支持 AMX 且任务适合矩阵 tile 的服务器 CPU。

然后在启动时选择最快的可用实现。

9. 一个简单的学习路线#

如果只是想理解 AVX/AMX，不建议一上来就手写 AMX GEMM。更合理的顺序是：

1. 先理解标量、SIMD、向量寄存器和 cache。
2. 用 C++ 写一个普通数组加法，让编译器自动向量化，配合 -O3 -mavx2 观察汇编。
3. 用 AVX intrinsics 写 _mm256_add_ps 这种小例子。
4. 学习矩阵乘法的分块：为什么要 block，为什么内存布局重要。
5. 再看 AMX tile 配置、tile load/store、tile dot-product。
6. 最后阅读 oneDNN、BLAS 或深度学习框架里的 CPU kernel 思路。

对 Python 用户来说，实用路线则是：

1. 确认 NumPy/PyTorch/TensorFlow 的构建和后端。
2. 用合适的数据类型，例如 BF16/INT8。
3. 关注模型形状和 batch size，让算子足够大，值得调用底层优化 kernel。
4. 用 profiler 验证，而不是凭感觉判断。

10. 总结#

AVX 和 AMX 都是 CPU 为了提高吞吐量而提供的指令集扩展，但它们服务的计算形态不同。

AVX 的核心是 SIMD：一条指令处理多个一维数据。它适合数组、向量、图像、信号处理和很多数值计算。C++ 可以通过自动向量化或 intrinsics 使用；Python 通常通过 NumPy、BLAS、PyTorch 等库间接使用。

AMX 的核心是 tile：把矩阵块放进二维 tile 寄存器，再用 TMUL 这类硬件单元做矩阵乘加。它更偏向 AI 训练和推理里的矩阵乘法，直接手写门槛较高，通常通过 oneDNN、MKL、PyTorch/TensorFlow CPU 后端间接使用。

最重要的一点是：硬件能力、操作系统支持、编译器/库支持必须同时成立。只看 CPU 名字或只看 Python 代码，都无法保证 AVX/AMX 真的被用上。

参考资料#

• Intel: Instruction Set Architecture
• Intel: Intel Intrinsics Guide
• Intel: Intel AVX-512 Overview
• Intel: What Is Intel Advanced Matrix Extensions?
• Intel: AMX intrinsics code sample