x86 和 AVX

Intel 和 AMD 的 x86 处理器主导着大多数 ML training 发生的数据中心服务器。该文件涵盖了 x86 SIMD、AVX/AVX2 intrinsics 编程的演变、用于矩阵运算的 AVX-512、Intel AMX、内存 alignment、性能陷阱和 profiling——用于从世界上最常见的服务器 CPU。

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

x86 SIMD 演变

• x86 SIMD 已由越来越宽的 vector registers 演变而来：

一代	年	寄存器宽度	寄存器	主要特点
MMX	1997	64位	8（毫米0-7）	仅整数，与 FPU 共享
SSE	1999	128位	8（xmm0-7）	4个浮子，专用registers
SSE2	2001	128位	8/16	2 个双精度数、整数运算
AVX	2011	256位	16（ymm0-15）	8 个浮点数，3 个操作数指令
AVX2	2013	256位	16	整数 256 位，FMA，聚集
AVX-512	2017	512位	32（zmm0-31）	16 个浮点数，掩码 registers，分散
AMX	2023	瓷砖 registers	8 block瓷砖	矩阵乘法（BF16、INT8）

• 每一代 vectorised 代码的吞吐量都增加了一倍。使用 SSE intrinsics 编写的代码可以在 2001 年以来制造的每个 x86 CPU 上运行。AVX2 需要 2013 年以上的 CPU。 AVX-512是Intel Xeon和一些消费类芯片。 AMX 是最新的（Sapphire Rapids 及更高版本）。

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

AVX2 编程

• AVX2 在 256 位 registers (YMM) 上运行，同时处理 8 个浮点型或 4 个双精度型。它是可移植高性能代码的最佳选择：几乎所有现代 x86 CPU（2013 年以上）都可用。

内在函数命名约定

• 所有 x86 intrinsics 均遵循模式：_mm[width]_[operation]_[type]

  • _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>  // all x86 SIMD intrinsics

// Data types
__m256  a;   // 256-bit register holding 8 float32s
__m256d b;   // 256-bit register holding 4 float64s
__m256i c;   // 256-bit register holding integers (8x32, 16x16, or 32x8)

加载和存储数据

// Load 8 floats from memory
__m256 v = _mm256_loadu_ps(ptr);      // unaligned load (works with any address)
__m256 v = _mm256_load_ps(ptr);       // aligned load (ptr must be 32-byte aligned, faster)

// Store 8 floats to memory
_mm256_storeu_ps(out_ptr, v);          // unaligned store
_mm256_store_ps(out_ptr, v);           // aligned store

// Broadcast a single value to all 8 lanes
__m256 ones = _mm256_set1_ps(1.0f);    // [1, 1, 1, 1, 1, 1, 1, 1]

// Set individual values (rarely needed)
__m256 v = _mm256_set_ps(7,6,5,4,3,2,1,0);  // note: reverse order!

// Zero register
__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] (slower than mul)

// Fused Multiply-Add: r = a * b + c (one instruction, one rounding)
__m256 r = _mm256_fmadd_ps(a, b, c);   // THE most important instruction for ML

// Min and max
__m256 mn = _mm256_min_ps(a, b);       // min(a[i], b[i]) — useful for clipping
__m256 mx = _mm256_max_ps(a, b);       // max(a[i], b[i]) — useful for ReLU

实例：AVX2 点积

#include <immintrin.h>

float dot_avx2(const float* a, const float* b, int n) {
    __m256 sum = _mm256_setzero_ps();  // 8 accumulators initialised to 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
    }

    // Horizontal reduction: sum the 8 elements of sum
    // Step 1: add upper 128 bits to lower 128 bits
    __m128 hi = _mm256_extractf128_ps(sum, 1);
    __m128 lo = _mm256_castps256_ps128(sum);
    __m128 sum128 = _mm_add_ps(hi, lo);        // 4 partial sums

    // Step 2: horizontal add within 128-bit register
    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);       // extract scalar

    // Scalar cleanup
    for (; i < n; i++) {
        result += a[i] * b[i];
    }

    return result;
}

• 为什么水平缩小丑陋：SIMD是为垂直操作而设计的（车道0与车道0，车道1与车道1）。水平操作（跨通道求和）与硬件对抗。这就是为什么点积最后有尴尬的归约代码。 vectorised 循环是干净的；减少是样板。

• 性能：与 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);
    }

    // Reduce 8 maxes to 1
    __m128 hi = _mm256_extractf128_ps(max_vec, 1);
    __m128 lo = _mm256_castps256_ps128(max_vec);
    __m128 max128 = _mm_max_ps(hi, lo);

    // Shuffle and max to find the single maximum
    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 指令重新排列 register 内的元素。二进制常量 0b01001110 控制哪些元素去哪里。这称为 置换，它直接连接到置换矩阵（第 2 章）：对 SIMD 通道进行混洗相当于乘以置换矩阵的硬件。

AVX-512

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

__m512 a = _mm512_loadu_ps(ptr);                // load 16 floats
__m512 c = _mm512_fmadd_ps(a, b, c);            // 16 FMAs at once
float sum = _mm512_reduce_add_ps(a);             // built-in horizontal sum (no manual reduction!)

// Mask operations: operate on a subset of lanes
__mmask16 mask = _mm512_cmpgt_ps_mask(a, zero);  // which lanes are > 0?
__m512 relu = _mm512_maskz_mov_ps(mask, a);       // zero out negative lanes = ReLU

• 面具registers（__mmask16）是AVX-512最强大的功能。每个位控制通道是否参与操作。这取代了标量清理循环：最后一次迭代使用仅有效通道处于活动状态的掩码，无需单独的标量循环即可处理任何 vector 长度。

• AVX-512 频率限制：在许多 Intel CPU 上，使用 AVX-512 指令会导致 CPU 暂时降低其时钟频率（以保持在热限制内）。这意味着对于短突发来说，AVX-512 并不总是比 AVX2 更快——频率损失可能超过更宽的向量。对于持续的工作负载（如矩阵乘法），AVX-512 获胜。对于混合代码（一些 SIMD，一些标量），频率转换可能会造成伤害。

Intel AMX：矩阵乘法硬件

• AMX（高级矩阵扩展）添加了专用矩阵乘法单元。 AMX 不是 SIMD 向量，而是在图block上运行：2D blocks 数据（每行最多 16 行 × 64 字节）。

#include <immintrin.h>

// AMX tile multiply: C += A * B (in BF16)
// A is 16x32 BF16, B is 32x16 BF16, C is 16x16 FP32
_tile_loadd(0, a_ptr, stride_a);   // load tile 0 from A
_tile_loadd(1, b_ptr, stride_b);   // load tile 1 from B
_tile_dpbf16ps(2, 0, 1);           // tile 2 += tile 0 * tile 1 (BF16 matmul, FP32 accumulation)
_tile_stored(2, c_ptr, stride_c);  // store tile 2 to C

• AMX 在一条指令中执行完整的 16×32 × 32×16 矩阵乘法。这是一次数百个 FMA 操作，专门为主导 transformer inference（attention 分数计算，MLP 层）的小型矩阵乘法而设计。

• AMX 支持 BF16 (bfloat16) 和 INT8，与 ML inference 中使用的精度匹配。与AVX-512结合进行其他操作，配备AMX的CPU（Intel Sapphire Rapids、Emerald Rapids）可以与入门级GPU竞争地运行transformer inference。

内存对齐

• 对齐内存访问是当数据地址是vector register宽度的倍数时（SSE为16字节，AVX为32字节，AVX-512为64字节）。对齐访问在某些 CPU 上速度更快，并且是 _mm256_load_ps 所要求的（与 _mm256_loadu_ps 相对）。

// Allocate aligned memory
float* data = (float*)aligned_alloc(32, n * sizeof(float));  // 32-byte aligned for AVX

// C++ aligned allocation
#include <new>
float* data = new (std::align_val_t(32)) float[n];

// Alternatively, use the compiler attribute
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 registers 的高 128 位。现代 CPU (Skylake+) 没有这种惩罚。

• 注册压力：AVX2有16 YMM registers。如果您的 kernel 使用太多变量，编译器会将 registers 溢出到堆栈（内存），从而破坏性能。使用很少的实时变量保持内部循环简单。

• 数据依赖性：sum = _mm256_fmadd_ps(a, b, sum) 依赖于 sum：每次迭代必须等待前一个 FMA 完成（约 4-5 个周期延迟）。修复：使用多个独立的累加器并在最后进行归约：

// Single accumulator: limited by FMA latency (4-5 cycles)
__m256 sum = _mm256_setzero_ps();
for (...) {
    sum = _mm256_fmadd_ps(a, b, sum);  // each depends on previous
}

// Four accumulators: 4x throughput (hide latency)
__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);  // independent
    sum1 = _mm256_fmadd_ps(a1, b1, sum1);  // independent
    sum2 = _mm256_fmadd_ps(a2, b2, sum2);  // independent
    sum3 = _mm256_fmadd_ps(a3, b3, sum3);  // independent
}
sum0 = _mm256_add_ps(sum0, sum1);
sum2 = _mm256_add_ps(sum2, sum3);
sum0 = _mm256_add_ps(sum0, sum2);

• 这是循环展开以隐藏延迟。 CPU 可以连续发出 FMA，因为它们写入不同的 registers。这是数字代码中最有影响力的微观优化之一。

分析

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

# Linux perf (requires kernel support)
perf stat ./my_program                    # basic counters: cycles, instructions, IPC
perf stat -e cache-misses,cache-references ./my_program  # cache behaviour
perf record -g ./my_program && perf report              # call graph profiling

# Intel VTune (detailed x86 profiling)
vtune -collect hotspots -- ./my_program
vtune -collect memory-access -- ./my_program   # memory bandwidth analysis

• 要寻找什么：
  • IPC（每周期指令）：CPU 的使用效率如何。 IPC > 2 较好。 IPC < 1 表明内存停顿或 branch 错误预测。
  • 缓存未命中率：L1/L2 未命中率较高表明数据局部性较差。重构数据访问模式。
  • branch误预测率：> 5% 表明不可预测的 branches。如果可能的话，转换为无branch代码（SIMD 比较+混合）。
  • 实现的 FLOPS 与最高线：将测量的 FLOPS 与最高线 model（文件 01）进行比较。如果你的水平低于屋顶线，那么还有改进的空间。

编码任务（在 x86 上使用 g++ 或 clang++ 进行编译 - Intel/AMD）

1. 编写标量点积和 AVX2 点积。对两者进行基准测试并测量 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);
       }
       // Reduce: add upper 128 to lower 128, then horizontal add
       __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;
   }
   
   // Unrolled: process 32 floats per iteration (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);  // warm up
           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;  // read + write
           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. 测量内存alignment的效果。比较大型阵列上的对齐负载和未对齐负载。

   // 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;
   
       // Aligned allocation (32-byte for AVX2)
       float* aligned = (float*)aligned_alloc(32, N * sizeof(float));
   
       // Unaligned: offset by 4 bytes (1 float) from aligned boundary
       float* raw = (float*)malloc((N + 1) * sizeof(float));
       float* unaligned = raw + 1;  // guaranteed misaligned
   
       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();
           // Warm up
           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;
   }