第3章 SIMD 指令集入门

第3章：SIMD 指令集入门：SSE/AVX 基础

说实话，很多搞C++高性能计算的朋友，一开始听到SIMD就头大。我当年也一样，觉得这东西太底层了，跟寄存器打交道，动不动就128位、256位，看着就晕。但后来我在做SwiGLU激活函数优化时，发现不用SIMD，性能根本提不上去。说白了，SIMD就是让你一条指令同时处理多个数据，这玩意儿是CPU给你开的「并行快车道」。

3.1 为什么需要SIMD？

你想想看，我们写一个循环：

for (int i = 0; i < 8; i++) {
    c[i] = a[i] + b[i];
}

普通情况下，CPU得一条一条地加，加8次。但如果你用AVX，一条指令就能把8个float一口气加完。嗯，这就是SIMD的威力——单指令多数据。

我在项目中遇到过这样一个场景：SwiGLU内核里有个逐元素乘加操作，数据量不大但调用频率极高。如果用普通循环，延迟全耗在指令发射上了。换成SIMD后，吞吐量直接翻了三倍多。所以，搞高性能计算，SIMD是绕不开的一关。

3.2 SSE vs AVX：选哪个？

SSE（Streaming SIMD Extensions）和AVX（Advanced Vector Extensions）是Intel/AMD CPU上最常见的两套SIMD指令集。我整理了个对比表，方便你快速了解：

特性	SSE	AVX	AVX-512
寄存器宽度	128位	256位	512位
单精度float数量	4个	8个	16个
双精度double数量	2个	4个	8个
指令前缀	_mm_	_mm256_	_mm512_
头文件	<xmmintrin.h>	<immintrin.h>	<immintrin.h>

我个人习惯是：能用AVX就别用SSE。为什么呢？因为AVX寄存器更宽，一次处理的数据更多。但要注意，老一点的CPU（比如2011年之前的）不支持AVX。如果你要写通用库，得做运行时检测。

⚠️ 避坑指南： 我曾经在服务器上部署代码时，忘了检查CPU是否支持AVX，结果程序直接崩溃——非法指令异常。后来我加了个cpuid检测，才彻底解决。记住：不要假设所有CPU都支持AVX。

3.3 数据类型：m128 / m256 到底是什么？

这些类型说白了就是「寄存器的C++化身」。你没法直接操作寄存器，但编译器给你提供了这些类型，让你能往寄存器里塞数据。

__m128：128位寄存器，存4个float
__m128d：128位寄存器，存2个double
__m128i：128位寄存器，存整数（8个short或4个int等）
__m256：256位寄存器，存8个float
__m256d：256位寄存器，存4个double
__m256i：256位寄存器，存整数

注意，这些类型不是普通的结构体，它们有特殊的内存对齐要求。__m128需要16字节对齐，__m256需要32字节对齐。如果你用普通malloc分配内存，然后强转成__m256*，大概率会崩。

💡 小技巧： 用 _mm_malloc 或 aligned_alloc 来分配对齐内存。我一般用 _mm_malloc(size, 32) 来保证AVX数据对齐。

3.4 基本算术操作：加法、减法、乘法、除法

好，现在我们来点实际的。假设你有两个数组，每个数组有8个float，你想把它们对应相加。用AVX怎么写？

#include <immintrin.h>

void add_arrays_avx(float* a, float* b, float* c) {
    // 从内存加载数据到寄存器
    __m256 va = _mm256_load_ps(a);   // 加载8个float
    __m256 vb = _mm256_load_ps(b);
    
    // 执行加法
    __m256 vc = _mm256_add_ps(va, vb);
    
    // 结果写回内存
    _mm256_store_ps(c, vc);
}

你看，核心就三行：加载、计算、存储。其他算术操作也类似：

操作	SSE (128位)	AVX (256位)
加法	_mm_add_ps	_mm256_add_ps
减法	_mm_sub_ps	_mm256_sub_ps
乘法	_mm_mul_ps	_mm256_mul_ps
除法	_mm_div_ps	_mm256_div_ps
乘加	_mm_add_ps(_mm_mul_ps(a,b), c)	_mm256_fmadd_ps

这里有个坑：除法在SIMD里很慢。我做过测试，_mm256_div_ps的延迟大约是加法的10倍。所以在SwiGLU里，我尽量用乘法代替除法，比如用乘以倒数来代替。

3.5 实战：用SIMD实现向量点积

点积是SwiGLU里经常出现的操作。我们来看一个AVX实现：

float dot_product_avx(float* a, float* b, int n) {
    __m256 sum = _mm256_setzero_ps();  // 初始化为0
    
    for (int i = 0; i < n; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);  // 非对齐加载
        __m256 vb = _mm256_loadu_ps(b + i);
        __m256 vprod = _mm256_mul_ps(va, vb);
        sum = _mm256_add_ps(sum, vprod);
    }
    
    // 将8个部分和归约为1个
    float result[8];
    _mm256_storeu_ps(result, sum);
    float final_sum = 0.0f;
    for (int i = 0; i < 8; i++) {
        final_sum += result[i];
    }
    return final_sum;
}

注意我用了_mm256_loadu_ps而不是_mm256_load_ps。区别在于：load要求32字节对齐，loadu不要求。如果你不确定数据是否对齐，用loadu更安全，但性能会稍微差一点。

    🔑 核心要点： SIMD编程的本质就是「数据并行」。你把数据打包进寄存器，然后一条指令搞定一批数据的操作。在SwiGLU里，我们经常用SIMD来加速逐元素的激活函数计算。

3.6 知识体系总览

下面这张图帮你理清本章的知识脉络：

嗯，这张图把本章的核心脉络串起来了。你从「为什么需要SIMD」出发，理解它的价值；然后对比SSE和AVX，选对工具；接着掌握数据类型和基本操作；最后通过实战把知识落地。说白了，SIMD没那么神秘，它就是让你用更少的指令干更多的活。

💡 我的建议： 刚开始学SIMD时，别急着写复杂逻辑。先写一个简单的加法，跑通，然后看汇编代码确认是否生成了SIMD指令。我在调试SwiGLU时就是这么一步步来的。

好了，这一章的内容就到这里。记住：SIMD是高性能计算的基石，SwiGLU的优化离不开它。下一章我们会深入SwiGLU的具体实现，到时候你会看到这些指令是怎么派上用场的。

📚 课程目录