maths-cs-ai-compendium-zh/chapter 16: SIMD and GPU programming/03. x86 and AVX.md

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