x86 与 AVX

一句话总结: x86 是数据中心的绝对霸主，掌握它的高级向量扩展 (Advanced Vector Extensions, AVX) / AVX2 / AVX-512 与高级矩阵扩展 (Advanced Matrix eXtensions, AMX) 内建函数，就能榨干服务器 CPU 的每一分性能。

来自 Intel 与 AMD 的 x86 处理器统治着数据中心服务器，大多数 ML 训练都在这上面跑。 本节覆盖 x86 单指令多数据 (Single Instruction Multiple Data, SIMD) 的演进、AVX / AVX2 内建函数编程、AVX-512、用于矩阵运算的 Intel AMX、内存对齐、性能陷阱与性能剖析 —— 这些是从全世界最常见的服务器 CPU 上榨出最大性能的全套工具.

• 如果你的训练跑在云虚拟机上 (AWS、GCP、Azure)，那几乎必然是 x86. 就算 GPU 占了大头，中央处理器 (Central Processing Unit, CPU) 侧也有瓶颈：数据加载、预处理、梯度聚合、检查点保存，全在 CPU 上。 用 x86 SIMD 优化这些环节，能实打实地压缩端到端训练时间。

x86 SIMD 演进史

• x86 SIMD 一路走来，向量寄存器越来越宽:

世代	年份	寄存器宽度	寄存器数量	关键特性
MMX	1997	64 位	8 (mm0-7)	仅整数，与浮点单元 (FPU) 共用
SSE	1999	128 位	8 (xmm0-7)	4 个 float，独立寄存器
SSE2	2001	128 位	8/16	2 个 double，整数运算
AVX	2011	256 位	16 (ymm0-15)	8 个 float，三操作数指令
AVX2	2013	256 位	16	256 位整数、融合乘加 (Fused Multiply-Add, FMA)、gather
AVX-512	2017	512 位	32 (zmm0-31)	16 个 float、掩码寄存器、scatter
AMX	2023	Tile 寄存器	8 个 tile	矩阵乘法 (BF16、INT8)

• 每一代都把向量化代码的吞吐量翻了一倍。 用 SSE 内建函数写的代码，在 2001 年之后出厂的任何 x86 CPU 上都能跑。 AVX2 要求 2013 年之后的 CPU. AVX-512 主要出现在 Intel Xeon 和部分消费级芯片上。 AMX 是最新的一代 (Sapphire Rapids 及之后).

• 向后兼容: x86 SSE 寄存器 (xmm) 就是 AVX 寄存器 (ymm) 的低 128 位，而 AVX 寄存器又是 AVX-512 寄存器 (zmm) 的低 256 位。 旧的 SSE 代码不改一行就能在新 CPU 上跑。

AVX2 编程

• AVX2 在 256 位寄存器 (YMM) 上工作，一次处理 8 个 float 或 4 个 double. 它是"可移植高性能代码"的甜蜜点：几乎所有现代 x86 CPU (2013 年之后) 都支持。

内建函数命名约定

• 所有 x86 内建函数都遵循同一套模式：_mm[宽度]_[操作]_[类型]

  • _mm = MMX/SSE (128 位), _mm256 = AVX (256 位), _mm512 = AVX-512 (512 位)
  • 操作：add、mul、fmadd、load、store、set 等等
  • 类型：ps = packed single (float32), pd = packed double (float64), epi32 = packed int32, si256 = 256 位整数

#include <immintrin.h>  // 所有 x86 SIMD 内建函数

// 数据类型
__m256  a;   // 256 位寄存器, 装 8 个 float32
__m256d b;   // 256 位寄存器, 装 4 个 float64
__m256i c;   // 256 位寄存器, 装整数 (8x32、16x16 或 32x8)

加载与存储数据

// 从内存加载 8 个 float
__m256 v = _mm256_loadu_ps(ptr);      // 非对齐加载 (任意地址都行)
__m256 v = _mm256_load_ps(ptr);       // 对齐加载 (ptr 必须 32 字节对齐, 更快)

// 存 8 个 float 到内存
_mm256_storeu_ps(out_ptr, v);          // 非对齐存储
_mm256_store_ps(out_ptr, v);           // 对齐存储

// 把一个值广播到全部 8 个通道
__m256 ones = _mm256_set1_ps(1.0f);    // [1, 1, 1, 1, 1, 1, 1, 1]

// 逐个设值 (极少用)
__m256 v = _mm256_set_ps(7,6,5,4,3,2,1,0);  // 注意: 顺序是反的!

// 清零寄存器
__m256 z = _mm256_setzero_ps();

算术运算

__m256 c = _mm256_add_ps(a, b);        // c[i] = a[i] + b[i]
__m256 d = _mm256_mul_ps(a, b);        // d[i] = a[i] * b[i]
__m256 e = _mm256_sub_ps(a, b);        // e[i] = a[i] - b[i]
__m256 f = _mm256_div_ps(a, b);        // f[i] = a[i] / b[i] (比 mul 慢)

// 融合乘加: r = a * b + c (一条指令, 一次舍入)
__m256 r = _mm256_fmadd_ps(a, b, c);   // ML 里最重要的一条指令

// 最小值与最大值
__m256 mn = _mm256_min_ps(a, b);       // min(a[i], b[i]) —— 适合做裁剪
__m256 mx = _mm256_max_ps(a, b);       // max(a[i], b[i]) —— 适合做 ReLU

实战示例：AVX2 点积

#include <immintrin.h>

float dot_avx2(const float* a, const float* b, int n) {
    __m256 sum = _mm256_setzero_ps();  // 8 个累加器初始化为 0

    int i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        sum = _mm256_fmadd_ps(va, vb, sum);  // sum += va * vb
    }

    // 水平规约: 把 sum 的 8 个元素加成 1 个
    // Step 1: 高 128 位加到低 128 位
    __m128 hi = _mm256_extractf128_ps(sum, 1);
    __m128 lo = _mm256_castps256_ps128(sum);
    __m128 sum128 = _mm_add_ps(hi, lo);        // 4 个部分和

    // Step 2: 在 128 位寄存器内做水平加
    sum128 = _mm_hadd_ps(sum128, sum128);       // [a+b, c+d, a+b, c+d]
    sum128 = _mm_hadd_ps(sum128, sum128);       // [a+b+c+d, ...]

    float result = _mm_cvtss_f32(sum128);       // 取出标量

    // 标量收尾
    for (; i < n; i++) {
        result += a[i] * b[i];
    }

    return result;
}

• 为什么水平规约这么丑: SIMD 天生就是给"垂直"操作准备的 (通道 0 配通道 0，通道 1 配通道 1). 水平操作 (跨通道求和) 是在跟硬件对着干。 所以点积末尾总有一段别扭的规约代码。 向量化的主循环很干净，规约部分则是模板化的样板。

• 性能：相比 NEON 版本 (第 02 节), AVX2 每次迭代处理 8 个 float，而 NEON 只能处理 4 个。 对长向量来说，这就是相对 NEON 的 2 倍加速 (先不算内存带宽的限制).

实战示例：AVX2 Softmax (简化版)

• Softmax 需要：找最大值、减掉最大值、求指数、求和、相除。 这里演示第一步 —— 找最大值:

float vector_max_avx2(const float* data, int n) {
    __m256 max_vec = _mm256_set1_ps(-INFINITY);

    int i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 v = _mm256_loadu_ps(data + i);
        max_vec = _mm256_max_ps(max_vec, v);
    }

    // 把 8 个 max 规约成 1 个
    __m128 hi = _mm256_extractf128_ps(max_vec, 1);
    __m128 lo = _mm256_castps256_ps128(max_vec);
    __m128 max128 = _mm_max_ps(hi, lo);

    // Shuffle + max 找出唯一的最大值
    max128 = _mm_max_ps(max128, _mm_shuffle_ps(max128, max128, 0b01001110));
    max128 = _mm_max_ps(max128, _mm_shuffle_ps(max128, max128, 0b10110001));

    float result = _mm_cvtss_f32(max128);

    for (; i < n; i++) {
        result = result > data[i] ? result : data[i];
    }

    return result;
}

• _mm_shuffle_ps 指令把寄存器内的元素重新排列。 二进制常量 0b01001110 控制哪个元素去哪。 这叫做置换 (permutation)，它直接和置换矩阵挂钩 (见第 2 章): shuffle SIMD 通道，就是"乘一个置换矩阵"的硬件版本。

AVX-512

• AVX-512 把宽度再翻一番：512 位寄存器 (ZMM)，一次处理 16 个 float.

__m512 a = _mm512_loadu_ps(ptr);                // 加载 16 个 float
__m512 c = _mm512_fmadd_ps(a, b, c);            // 一次 16 个 FMA
float sum = _mm512_reduce_add_ps(a);             // 内建的水平求和 (不用手写规约了!)

// 掩码操作: 只对部分通道生效
__mmask16 mask = _mm512_cmpgt_ps_mask(a, zero);  // 哪些通道 > 0?
__m512 relu = _mm512_maskz_mov_ps(mask, a);       // 把负数通道清零 = ReLU

• 掩码寄存器 (__mmask16) 是 AVX-512 最强大的特性。 每一位控制一条通道是否参与运算。 这直接干掉了那段标量收尾循环：最后一次迭代只要用一个"只有有效通道被激活"的掩码，就能处理任意长度的向量，不需要单独写标量循环。

• AVX-512 频率降档：很多 Intel CPU 上，一旦用了 AVX-512 指令，CPU 会暂时降低时钟频率 (为了压住热量). 也就是说，AVX-512 未必总是比 AVX2 快 —— 短时突发场景下，降频损失可能盖过更宽向量带来的收益。 对持续型负载 (比如矩阵乘法), AVX-512 稳赢。 对混合代码 (一段 SIMD、一段标量)，频率来回切换反而会伤性能。

Intel AMX：矩阵乘法专用硬件

• AMX (Advanced Matrix eXtensions) 加入了专用矩阵乘法单元。 它不再操作 SIMD 向量，而是操作 tile：二维数据块 (最多 16 行 × 每行 64 字节).

#include <immintrin.h>

// AMX tile 乘法: C += A * B (BF16 精度)
// A 是 16x32 BF16, B 是 32x16 BF16, C 是 16x16 FP32
_tile_loadd(0, a_ptr, stride_a);   // 从 A 加载到 tile 0
_tile_loadd(1, b_ptr, stride_b);   // 从 B 加载到 tile 1
_tile_dpbf16ps(2, 0, 1);           // tile 2 += tile 0 * tile 1 (BF16 矩阵乘, FP32 累加)
_tile_stored(2, c_ptr, stride_c);  // tile 2 存回 C

• AMX 用一条指令就完成一次完整的 16×32 × 32×16 矩阵乘法。 这相当于同时做几百个 FMA 操作，专门为 Transformer 推理里那些小矩阵乘法量身设计 (注意力分数计算、MLP 层).

• AMX 支持 BF16 (bfloat16) 与 INT8，精度选型恰好匹配 ML 推理。 再叠加 AVX-512 处理其他运算，装了 AMX 的 CPU (Intel Sapphire Rapids、Emerald Rapids) 跑 Transformer 推理，可以在性能上与入门级 GPU 掰手腕。

内存对齐

• 对齐内存访问 指的是数据地址是向量寄存器宽度的整数倍 (SSE 是 16 字节、AVX 是 32 字节、AVX-512 是 64 字节). 在部分 CPU 上对齐访问更快，而且 _mm256_load_ps 强制要求对齐 (对应地，_mm256_loadu_ps 不要求).

// 分配对齐内存
float* data = (float*)aligned_alloc(32, n * sizeof(float));  // AVX 要 32 字节对齐

// C++ 对齐分配
#include <new>
float* data = new (std::align_val_t(32)) float[n];

// 或者用编译器属性
alignas(32) float data[1024];

• 实际情况：在现代 CPU (Haswell 及之后) 上，只要数据不跨越缓存行边界，非对齐加载 (loadu) 几乎和对齐加载一样快。 非对齐访问的性能惩罚基本已经消失，但缓存行分裂 (数据跨越两条 64 字节缓存行) 仍会给那次加载带来约 2 倍的减速。 用对齐分配可以彻底避掉这个问题。

性能陷阱

• AVX-SSE 切换惩罚：在老一代 Intel CPU (Skylake 之前) 上，在 AVX (256 位) 与 SSE (128 位) 指令之间切换会带来性能损失 (约 70 周期). 这就是为什么在用 AVX 的函数返回前，你应该调用 _mm256_zeroupper() (或直接执行 vzeroupper 指令)，把 YMM 寄存器的高 128 位清掉。 现代 CPU (Skylake+) 已经没有这个惩罚了。

• 寄存器压力: AVX2 只有 16 个 YMM 寄存器。 如果你的内核用了太多变量，编译器就会把寄存器溢出 (spill) 到栈上 (内存)，性能直接崩。 内层循环要保持简单，活跃变量越少越好。

• 数据依赖: sum = _mm256_fmadd_ps(a, b, sum) 依赖于 sum —— 每次迭代都得等上一次 FMA 算完 (延迟约 4-5 周期). 解法：用多个独立累加器，最后再规约:

// 单累加器: 受限于 FMA 延迟 (4-5 周期)
__m256 sum = _mm256_setzero_ps();
for (...) {
    sum = _mm256_fmadd_ps(a, b, sum);  // 每次都依赖上一次
}

// 4 个累加器: 吞吐 4 倍 (掩盖延迟)
__m256 sum0 = _mm256_setzero_ps();
__m256 sum1 = _mm256_setzero_ps();
__m256 sum2 = _mm256_setzero_ps();
__m256 sum3 = _mm256_setzero_ps();
for (...) {
    sum0 = _mm256_fmadd_ps(a0, b0, sum0);  // 相互独立
    sum1 = _mm256_fmadd_ps(a1, b1, sum1);  // 相互独立
    sum2 = _mm256_fmadd_ps(a2, b2, sum2);  // 相互独立
    sum3 = _mm256_fmadd_ps(a3, b3, sum3);  // 相互独立
}
sum0 = _mm256_add_ps(sum0, sum1);
sum2 = _mm256_add_ps(sum2, sum3);
sum0 = _mm256_add_ps(sum0, sum2);

• 这就是循环展开 (loop unrolling)，用来掩盖延迟。 CPU 可以一条接一条把 FMA 发出去，因为它们写到不同的寄存器。 这是数值代码里最有效的微优化之一。

性能剖析

• 性能计数器 提供硬件级测量数据:

# Linux perf (需要内核支持)
perf stat ./my_program                    # 基础计数器: cycles、instructions、IPC
perf stat -e cache-misses,cache-references ./my_program  # 缓存行为
perf record -g ./my_program && perf report              # 调用图剖析

# Intel VTune (x86 详细剖析)
vtune -collect hotspots -- ./my_program
vtune -collect memory-access -- ./my_program   # 内存带宽分析

• 该盯什么:
  • 每周期指令数 (Instructions Per Cycle, IPC)：衡量 CPU 用得有多充分。 IPC > 2 算不错。 IPC < 1 通常暗示内存停顿或分支预测失败。
  • 缓存缺失率: L1 / L2 缺失率高说明数据局部性差，重新组织访问模式。
  • 分支预测失败率: > 5% 说明分支不可预测。 尽量改成无分支代码 (SIMD 比较 + 混合).
  • 实测 FLOPS 对比 Roofline：把测出的 FLOPS 和 Roofline 模型 (第 01 节) 对比，如果低于屋顶线，就还有提升空间。

编码任务 (在 x86 —— Intel/AMD 上用 g++ 或 clang++ 编译)

1. 写一个标量点积和一个 AVX2 点积，分别 benchmark，测出 8 宽 SIMD 带来的加速比。

   // task1_avx_dot.cpp
   // Compile: g++ -O3 -mavx2 -mfma -o task1 task1_avx_dot.cpp
   
   #include <iostream>
   #include <chrono>
   #include <vector>
   #include <immintrin.h>
   
   float dot_scalar(const float* a, const float* b, int n) {
       float sum = 0.0f;
       for (int i = 0; i < n; i++) sum += a[i] * b[i];
       return sum;
   }
   
   float dot_avx2(const float* a, const float* b, int n) {
       __m256 sum = _mm256_setzero_ps();
       int i = 0;
       for (; i + 8 <= n; i += 8) {
           __m256 va = _mm256_loadu_ps(a + i);
           __m256 vb = _mm256_loadu_ps(b + i);
           sum = _mm256_fmadd_ps(va, vb, sum);
       }
       // 规约: 高 128 位加到低 128 位, 再做水平加
       __m128 hi = _mm256_extractf128_ps(sum, 1);
       __m128 lo = _mm256_castps256_ps128(sum);
       __m128 r = _mm_add_ps(hi, lo);
       r = _mm_hadd_ps(r, r);
       r = _mm_hadd_ps(r, r);
       float result = _mm_cvtss_f32(r);
       for (; i < n; i++) result += a[i] * b[i];
       return result;
   }
   
   int main() {
       const int N = 10'000'000;
       std::vector<float> a(N, 1.0f), b(N, 2.0f);
   
       volatile float s1 = dot_scalar(a.data(), b.data(), N);
       volatile float s2 = dot_avx2(a.data(), b.data(), N);
   
       auto bench = [&](auto fn, const char* name) {
           auto start = std::chrono::high_resolution_clock::now();
           volatile float s;
           for (int t = 0; t < 100; t++) s = fn(a.data(), b.data(), N);
           auto end = std::chrono::high_resolution_clock::now();
           double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
           std::cout << name << ": " << ms << " ms (result: " << s << ")\n";
           return ms;
       };
   
       double t1 = bench(dot_scalar, "Scalar");
       double t2 = bench(dot_avx2,   "AVX2  ");
       std::cout << "Speedup: " << t1 / t2 << "x\n";
       return 0;
   }
   

2. 用 _mm256_max_ps 实现 AVX2 版本的 ReLU，与标量循环对比。 再试试多累加器 (循环展开) 掩盖 FMA 延迟。

   // task2_avx_relu.cpp
   // Compile: g++ -O3 -mavx2 -o task2 task2_avx_relu.cpp
   
   #include <iostream>
   #include <chrono>
   #include <vector>
   #include <immintrin.h>
   
   void relu_scalar(const float* in, float* out, int n) {
       for (int i = 0; i < n; i++) {
           out[i] = in[i] > 0.0f ? in[i] : 0.0f;
       }
   }
   
   void relu_avx2(const float* in, float* out, int n) {
       __m256 zero = _mm256_setzero_ps();
       int i = 0;
       for (; i + 8 <= n; i += 8) {
           __m256 x = _mm256_loadu_ps(in + i);
           _mm256_storeu_ps(out + i, _mm256_max_ps(x, zero));
       }
       for (; i < n; i++) out[i] = in[i] > 0.0f ? in[i] : 0.0f;
   }
   
   // 展开版: 每次迭代处理 32 个 float (4 x 8)
   void relu_avx2_unrolled(const float* in, float* out, int n) {
       __m256 zero = _mm256_setzero_ps();
       int i = 0;
       for (; i + 32 <= n; i += 32) {
           __m256 x0 = _mm256_loadu_ps(in + i);
           __m256 x1 = _mm256_loadu_ps(in + i + 8);
           __m256 x2 = _mm256_loadu_ps(in + i + 16);
           __m256 x3 = _mm256_loadu_ps(in + i + 24);
           _mm256_storeu_ps(out + i,      _mm256_max_ps(x0, zero));
           _mm256_storeu_ps(out + i + 8,  _mm256_max_ps(x1, zero));
           _mm256_storeu_ps(out + i + 16, _mm256_max_ps(x2, zero));
           _mm256_storeu_ps(out + i + 24, _mm256_max_ps(x3, zero));
       }
       for (; i + 8 <= n; i += 8) {
           _mm256_storeu_ps(out + i, _mm256_max_ps(_mm256_loadu_ps(in + i), zero));
       }
       for (; i < n; i++) out[i] = in[i] > 0.0f ? in[i] : 0.0f;
   }
   
   int main() {
       const int N = 16'000'000;
       std::vector<float> in(N), out(N);
       for (int i = 0; i < N; i++) in[i] = (float)(i % 200) - 100.0f;
   
       auto bench = [&](auto fn, const char* name) {
           fn(in.data(), out.data(), N);  // 预热
           auto start = std::chrono::high_resolution_clock::now();
           for (int t = 0; t < 100; t++) fn(in.data(), out.data(), N);
           auto end = std::chrono::high_resolution_clock::now();
           double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
           double bw = 2.0 * N * sizeof(float) / ms / 1e6;  // 读 + 写
           std::cout << name << ": " << ms << " ms (" << bw << " GB/s)\n";
       };
   
       bench(relu_scalar,        "Scalar        ");
       bench(relu_avx2,          "AVX2          ");
       bench(relu_avx2_unrolled, "AVX2 unrolled ");
       return 0;
   }
   

3. 测量内存对齐的效果，在大数组上对比对齐加载与非对齐加载。

   // task3_alignment.cpp
   // Compile: g++ -O3 -mavx2 -o task3 task3_alignment.cpp
   
   #include <iostream>
   #include <chrono>
   #include <cstdlib>
   #include <immintrin.h>
   
   int main() {
       const int N = 16'000'000;
   
       // 对齐分配 (AVX2 要 32 字节)
       float* aligned = (float*)aligned_alloc(32, N * sizeof(float));
   
       // 非对齐: 从对齐边界偏移 4 字节 (1 个 float)
       float* raw = (float*)malloc((N + 1) * sizeof(float));
       float* unaligned = raw + 1;  // 保证错位
   
       for (int i = 0; i < N; i++) {
           aligned[i] = 1.0f;
           unaligned[i] = 1.0f;
       }
   
       auto bench = [&](float* ptr, bool use_aligned, const char* name) {
           __m256 sum = _mm256_setzero_ps();
           // 预热
           for (int i = 0; i + 8 <= N; i += 8) {
               __m256 v = use_aligned ? _mm256_load_ps(ptr + i) : _mm256_loadu_ps(ptr + i);
               sum = _mm256_add_ps(sum, v);
           }
   
           auto start = std::chrono::high_resolution_clock::now();
           for (int t = 0; t < 100; t++) {
               sum = _mm256_setzero_ps();
               for (int i = 0; i + 8 <= N; i += 8) {
                   __m256 v = use_aligned ? _mm256_load_ps(ptr + i) : _mm256_loadu_ps(ptr + i);
                   sum = _mm256_add_ps(sum, v);
               }
           }
           auto end = std::chrono::high_resolution_clock::now();
           double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
           double bw = (double)N * sizeof(float) / ms / 1e6;
           std::cout << name << ": " << ms << " ms (" << bw << " GB/s)\n";
       };
   
       bench(aligned,   true,  "Aligned load  ");
       bench(unaligned, false, "Unaligned load");
   
       free(aligned);
       free(raw);
       return 0;
   }
   

---

📌 本节要点

• x86 SIMD 演进: MMX → SSE → AVX → AVX2 → AVX-512 → AMX，寄存器宽度从 64 位一路翻到 512 位，且完全向后兼容。
• AVX2 是甜蜜点: 256 位、一次 8 个 float, 2013 年之后的 x86 全支持，可移植高性能代码首选它。
• 内建函数命名: _mm[宽度]_[操作]_[类型], ps = float32、pd = float64、epi32 = int32，一看名字就知道操作粒度。
• FMA 是核心: _mm256_fmadd_ps 一条指令完成乘加、只做一次舍入，是 ML 数值代码里最重要的指令。
• 水平规约丑但必要：点积末尾要跨通道求和，与 SIMD 的垂直设计相悖，所以尾部代码总是又长又琐碎。
• AVX-512 掩码 + 降频权衡：掩码寄存器一次搞定标量收尾；但触发 AVX-512 可能导致 CPU 降频，短突发未必比 AVX2 快。
• Intel AMX 打矩阵乘：用 tile 一条指令完成 16×32 × 32×16 的 BF16 矩阵乘法，让 CPU 也能跟入门 GPU 掰手腕。
• 对齐 vs 非对齐：现代 CPU 上非对齐加载几乎无损，除非跨越缓存行，建议仍用 aligned_alloc(32, ...) 分配。
• 多累加器掩盖延迟: FMA 有 4-5 周期依赖延迟，展开成 4 个独立累加器可获得 4 倍吞吐 —— 数值代码里最有效的微优化之一。
• 性能剖析看 IPC + 缓存 + 分支: perf stat / VTune 抓 IPC、缓存缺失率、分支预测失败率，对比 Roofline 模型判断是否还有优化空间。