# 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位整数 ```cpp #include // 所有x86 SIMD内联函数 // 数据类型 __m256 a; // 256位寄存器，保存8个float32 __m256d b; // 256位寄存器，保存4个float64 __m256i c; // 256位寄存器，保存整数（8x32、16x16或32x8） ``` ### 加载和存储数据 ```cpp // 从内存加载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(); ``` ### 算术运算 ```cpp __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点积 ```cpp #include 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需要：找到最大值，减去它，求指数，求和，除法。以下是最值查找步骤： ```cpp 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个浮点数。 ```cpp __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字节）。 ```cpp #include // 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`）的要求。 ```cpp // 分配对齐内存 float* data = (float*)aligned_alloc(32, n * sizeof(float)); // AVX的32字节对齐 // C++对齐分配 #include 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个周期的延迟）。解决方案：使用多个独立累加器并在结束时归约： ```cpp // 单累加器：受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，因为它们写入不同的寄存器。这是数值代码中最有影响力的微优化之一。 ## 性能分析 - **性能计数器**提供硬件级测量： ```bash # 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带来的加速比。 ```cpp // task1_avx_dot.cpp // 编译：g++ -O3 -mavx2 -mfma -o task1 task1_avx_dot.cpp #include #include #include #include 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 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(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延迟。 ```cpp // task2_avx_relu.cpp // 编译：g++ -O3 -mavx2 -o task2 task2_avx_relu.cpp #include #include #include #include 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 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(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. 测量内存对齐的效果。比较在大数组上的对齐加载与非对齐加载。 ```cpp // task3_alignment.cpp // 编译：g++ -O3 -mavx2 -o task3 task3_alignment.cpp #include #include #include #include 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(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; } ```