x86 与 AVX

x86 与 AVX

来自 Intel 和 AMD 的 x86 处理器统治着数据中心服务器，而大多数 ML 训练都发生在那里。本节讲 x86 SIMD 的演进、AVX/AVX2 内联函数编程、AVX-512、用于矩阵运算的 Intel AMX、内存对齐、性能陷阱，以及性能剖析（profiling）——这些是从世界上最常见的服务器 CPU 上榨取最大性能的工具。

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

x86 SIMD 的演进

• x86 SIMD 通过越来越宽的向量寄存器不断演进：

• 每一代都让向量化代码的吞吐翻番。用 SSE 内联函数写的代码能跑在 2001 年以来生产的每一颗 x86 CPU 上。AVX2 需要 2013 年以后的 CPU。AVX-512 主要是 Intel Xeon 和部分消费级芯片。AMX 是最新的（Sapphire Rapids 及之后）。

• 向后兼容：x86 SSE 寄存器（xmm）是 AVX 寄存器（ymm）的低 128 位，而 ymm 又是 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;   // 持有 8 个 float32 的 256 位寄存器
__m256d b;   // 持有 4 个 float64 的 256 位寄存器
__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 步：把高 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 个 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 个最大值归约成 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 章的置换矩阵直接相关：对 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
​

• **掩码寄存器（mask register，__mmask16）**是 AVX-512 最强大的特性。每一位控制一个通道是否参与运算。这取代了标量收尾循环：最后一次迭代用一个掩码，只激活有效通道，无需单独的标量循环就能处理任意向量长度。

• AVX-512 的降频：在许多 Intel CPU 上，使用 AVX-512 指令会让 CPU 临时降低主频（以留在散热限制内）。这意味着对短促的工作，AVX-512 不一定比 AVX2 快——降频的代价可能超过更宽向量带来的收益。对持续的工作负载（如矩阵乘法），AVX-512 胜出；对混合代码（一部分 SIMD、一部分标量），频率切换可能反而拖累。

Intel AMX：矩阵乘硬件

• AMX（Advanced Matrix eXtensions，高级矩阵扩展）加入了专用的矩阵乘单元。AMX 不操作 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 matmul，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 一较高下。

内存对齐

• **对齐的内存访问（aligned memory access）**指数据地址是向量寄存器宽度的整数倍（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);
​

• 这就是用**循环展开（loop unrolling）**来隐藏延迟。CPU 能把 FMA 背靠背地发射，因为它们写到不同的寄存器。这是数值代码中最有影响力的微优化之一。

性能剖析（profiling）

• **性能计数器（performance counter）**提供硬件层面的测量：

# 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（Instructions Per Cycle，每周期指令数）：CPU 利用得有多高效。IPC > 2 不错；IPC < 1 暗示有内存停顿或分支误预测。
  • 缓存未命中率：高 L1/L2 未命中率说明数据局部性差，应重构访存模式。
  • 分支误预测率：> 5% 暗示有难以预测的分支，可能的话转成无分支代码（SIMD 比较 + 混合）。
  • 实际 FLOPS vs 屋顶线：把你测得的 FLOPS 与屋顶线模型（第 01 节）对比。如果你低于屋顶线，就还有提升空间。

编程练习（在 x86 上用 g++ 或 clang++ 编译——Intel/AMD）

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 (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
// 编译：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 × 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
// 编译：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;
}
​