跳转至

# x86与AVX

x86处理器来自Intel和AMD，主导着大多数ML训练所在的数据中心服务器。本文涵盖x86 SIMD的演进、AVX/AVX2内联函数编程、AVX-512、用于矩阵操作的Intel AMX、内存对齐、性能陷阱以及性能分析——在全球最常见的服务器CPU上榨取最大性能的工具。

- 如果你的训练在云虚拟机（AWS、GCP、Azure）上运行，它几乎肯定运行在x86上。即使是GPU密集训练也有CPU瓶颈：数据加载、预处理、梯度聚合和检查点保存都在CPU上运行。使用x86 SIMD优化这些环节可以有意义地减少端到端训练时间。

## x86 SIMD演进

- x86 SIMD经历了越来越宽的向量寄存器：

代次 | 年份 | 寄存器宽度 | 寄存器数量 | 关键特性
--- | --- | --- | --- | ---
MMX | 1997 | 64位 | 8（mm0-7） | 仅整数，与FPU共享
SSE | 1999 | 128位 | 8（xmm0-7） | 4个浮点数，专用寄存器
SSE2 | 2001 | 128位 | 8/16 | 2个双精度浮点数，整数操作
AVX | 2011 | 256位 | 16（ymm0-15） | 8个浮点数，三操作数指令
AVX2 | 2013 | 256位 | 16 | 整数256位，FMA，收集
AVX-512 | 2017 | 512位 | 32（zmm0-31） | 16个浮点数，掩码寄存器，分散
AMX | 2023 | 瓦片寄存器 | 8个瓦片 | 矩阵乘法（BF16，INT8）

- 每一代都将向量化代码的吞吐量翻倍。用SSE内联函数编写的代码可以在2001年以来制造的每一个x86 CPU上运行。AVX2需要2013年以后的CPU。AVX-512需要Intel Xeon和一些消费级芯片。AMX是最新的（Sapphire Rapids及以后）。

- 向后兼容性：x86 SSE寄存器（xmm）是AVX寄存器（ymm）的低128位，后者是AVX-512寄存器（zmm）的低256位。旧的SSE代码无需修改即可在新的CPU上运行。

## AVX2编程

- AVX2操作256位寄存器（YMM），同时处理8个浮点数或4个双精度浮点数。它是可移植高性能代码的甜点区域：在几乎所有现代x86 CPU（2013+）上可用。

### 内联函数命名约定

- 所有x86内联函数遵循模式：`_mm[宽度]_[操作]_[类型]`

  - `_mm` = MMX/SSE（128位），`_mm256` = AVX（256位），`_mm512` = AVX-512（512位）
  - 操作：`add`、`mul`、`fmadd`、`load`、`store`、`set` 等
  - 类型：`ps` = 打包单精度（float32），`pd` = 打包双精度（float64），`epi32` = 打包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个浮点数
    __m256 v = _mm256_loadu_ps(ptr);      // 非对齐加载（适用于任何地址）
    __m256 v = _mm256_load_ps(ptr);       // 对齐加载（ptr必须32字节对齐，更快）

    // 存储8个浮点数到内存
    _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]（比乘法慢）

    // 融合乘加：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：将上128位加到下128位
        __m128 hi = _mm256_extractf128_ps(sum, 1);
        __m128 lo = _mm256_castps256_ps128(sum);
        __m128 sum128 = _mm_add_ps(hi, lo);        // 4个部分和

        // 步骤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个浮点数，而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个最大值归约为1个
        __m128 hi = _mm256_extractf128_ps(max_vec, 1);
        __m128 lo = _mm256_castps256_ps128(max_vec);
        __m128 max128 = _mm_max_ps(hi, lo);

        // 通过混洗和取最大值找到单一最大值
        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` 控制哪些元素去哪里。这称为置换，它直接连接到置换矩阵（第2章）：打乱SIMD通道是做硬件级别的乘以置换矩阵。

## AVX-512

- AVX-512再次加倍宽度：512位寄存器（ZMM），同时处理16个浮点数。

    __m512 a = _mm512_loadu_ps(ptr);                // 加载16个浮点数
    __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（高级矩阵扩展）增加了专用矩阵乘法单元。AMX操作的不是SIMD向量，而是瓦片：2D数据块（最多16行 × 每行64字节）。

    #include <immintrin.h>

    // AMX瓦片乘法：C += A * B（BF16格式）
    // A为16x32 BF16，B为32x16 BF16，C为16x16 FP32
    _tile_loadd(0, a_ptr, stride_a);   // 从A加载瓦片0
    _tile_loadd(1, b_ptr, stride_b);   // 从B加载瓦片1
    _tile_dpbf16ps(2, 0, 1);           // 瓦片2 += 瓦片0 * 瓦片1（BF16矩阵乘法，FP32累加）
    _tile_stored(2, c_ptr, stride_c);  // 存储瓦片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寄存器。如果你的核函数使用太多变量，编译器会将寄存器溢出到栈（内存），从而破坏性能。保持内循环简单，活变量少。

- 数据依赖：`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倍吞吐量（隐藏延迟）
    __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);

- 这是循环展开以隐藏延迟。CPU可以背靠背发出FMAs，因为它们写入不同的寄存器。这是数值代码中最有影响力的微优化之一。

## 性能分析

- 性能计数器提供硬件级测量：

    # Linux perf（需要内核支持）
    perf stat ./my_program                    # 基本计数器：周期、指令、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   # 内存带宽分析

- 需要关注什么：
  - IPC（每周期指令数）：CPU被使用的效率。IPC > 2 良好。IPC < 1 表明内存停顿或分支预测错误。
  - 缓存缺失率：高L1/L2缺失率表明数据局部性差。需重构数据访问模式。
  - 分支预测错误率：> 5% 表明分支不可预测。如可能，转换为无分支代码（SIMD比较+混合）。
  - 实际FLOPS vs 屋顶线：将你的实测FLOPS与屋顶线模型（文件01）比较。如果你低于屋顶线，还有改进空间。

## 编程任务（在x86——Intel/AMD上用g++或clang++编译）

1. 编写标量点积和AVX2点积。对两者进行基准测试并测量8路SIMD带来的加速比。

       // task1_avx_dot.cpp
       // 编译：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（结果: " << s << "）\n";
               return ms;
           };

           double t1 = bench(dot_scalar, "标量");
           double t2 = bench(dot_avx2,   "AVX2  ");
           std::cout << "加速比: " << t1 / t2 << "x\n";
           return 0;
       }

2. 使用 `_mm256_max_ps` 实现AVX2 ReLU并与标量循环比较。然后尝试使用多累加器（循环展开）以隐藏FMA延迟。

       // task2_avx_relu.cpp
       // 编译：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个浮点数（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,        "标量          ");
           bench(relu_avx2,          "AVX2          ");
           bench(relu_avx2_unrolled, "AVX2展开      ");
           return 0;
       }

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

       // task3_alignment.cpp
       // 编译：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* 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,  "对齐加载  ");
           bench(unaligned, false, "非对齐加载");

           free(aligned);
           free(raw);
           return 0;
       }