「向量化与 SIMD」

// 标量版本: 每次处理一个元素
void add_scalar(float *restrict a, float *restrict b, float *restrict c, int n) {
    for (int i = 0; i < n; i++)
        c[i] = a[i] + b[i];
}

// 编译器在 -O2 -march=native 下会自动生成类似下面的向量版本
#include <immintrin.h>
void add_avx(float *restrict a, float *restrict b, float *restrict c, int n) {
    int i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        _mm256_storeu_ps(c + i, _mm256_add_ps(va, vb));
    }
    for (; i < n; i++) c[i] = a[i] + b[i];   // 处理尾巴
}

# 观察自动向量化的结果
cc -O3 -march=native -fopt-info-vec-optimized vec.c -c
cc -O3 -march=native -fopt-info-vec-missed vec.c -c     # 为什么没向量化
clang -O3 -march=native -Rpass=loop-vectorize vec.c -c
objdump -d vec.o | grep -E 'vaddps|vmovups' | head

// 存在真实依赖: 无法直接向量化
for (int i = 1; i < n; i++)
    a[i] = a[i - 1] + 1;          // 每次迭代依赖前一次的结果

// 依赖距离为 k 时, 可以按 k 个一组做向量化
for (int i = 4; i < n; i++)
    a[i] = a[i - 4] * 2;          // 距离 4, 分成 4 条独立链

// 反依赖: 写入的位置是下一次迭代要读的
for (int i = 0; i < n; i++)
    a[i] = b[i] + 1, b[i + 1] = a[i];   // 顺序敏感

// 别名是向量化最常见的阻碍
void scale(int *a, int *b, int n) {
    for (int i = 0; i < n; i++) a[i] = b[i] * 2;   // a 与 b 可能别名, 不向量化
}
void scale_fast(int *restrict a, int *restrict b, int n) {
    for (int i = 0; i < n; i++) a[i] = b[i] * 2;   // 承诺不重叠, 向量化
}

# 极简依赖判定: 仿射下标下的距离计算
def distance(dim_a, dim_b):
    """下标为 i*c + k 形式时, 距离 = (k_b - k_a) / (c_a - c_b)."""
    ca, ka = dim_a
    cb, kb = dim_b
    if ca == cb:
        return 0 if ka == kb else None      # 同址或平行不相交
    num, den = kb - ka, ca - cb
    return num // den if num % den == 0 else None

print(distance((1, 0), (1, -1)))   # a[i] 与 a[i-1]: 距离 1, 有依赖
print(distance((1, 0), (1, -4)))   # a[i] 与 a[i-4]: 距离 4
print(distance((2, 0), (2, 1)))    # a[2i] 与 a[2i+1]: 距离非整数, 无依赖

// 原始循环
for (int i = 0; i < n; i++) c[i] = a[i] * b[i] + 1.0f;

// 向量化后的等价形态 (VF = 8)
for (int i = 0; i + 8 <= n; i += 8) {
    __m256 va = _mm256_loadu_ps(a + i);
    __m256 vb = _mm256_loadu_ps(b + i);
    __m256 vc = _mm256_fmadd_ps(va, vb, _mm256_set1_ps(1.0f));  // FMA
    _mm256_storeu_ps(c + i, vc);
}
for (int i = n - n % 8; i < n; i++) c[i] = a[i] * b[i] + 1.0f;   // 尾巴

def vectorize_plan(n, vf, has_tail_check=True):
    """生成向量化的迭代划分方案."""
    plan = {"vector_iters": n // vf, "vf": vf, "tail": n % vf}
    if has_tail_check:
        plan["peel"] = 0
        plan["runtime_alias_check"] = True   # 无法静态证明时插入运行时检查
    return plan

print(vectorize_plan(1000, 8))
# {'vector_iters': 125, 'vf': 8, 'tail': 0, 'peel': 0, 'runtime_alias_check': True}

// 展开后的循环体: 4 条同构语句
c[0] = a[0] * b[0];
c[1] = a[1] * b[1];
c[2] = a[2] * b[2];
c[3] = a[3] * b[3];
// SLP 合并为: c[0..3] = a[0..3] * b[0..3]  (一条 128 位乘法)

def slp_pack(stmts, width):
    """把形状相同的标量语句按 width 打包."""
    groups, out = {}, []
    for s in stmts:
        key = (s[0], s[2])            # 操作类型 + 操作数形状
        groups.setdefault(key, []).append(s)
    for key, items in groups.items():
        while len(items) >= width:
            batch, items = items[:width], items[width:]
            out.append((f"vec{width}", key, [b[1] for b in batch]))
        out.extend(items)             # 不足一批的保持标量
    return out

stmts = [("mul", f"c{i}", ("a", "b", i)) for i in range(6)]
for op in slp_pack(stmts, 4):
    print(op)

// 标量归约
float sum = 0;
for (int i = 0; i < n; i++) sum += a[i];

// 向量化: 8 条独立部分和 + 末尾横向相加
__m256 vsum = _mm256_setzero_ps();
for (int i = 0; i + 8 <= n; i += 8)
    vsum = _mm256_add_ps(vsum, _mm256_loadu_ps(a + i));
float lanes[8];
_mm256_storeu_ps(lanes, vsum);
for (int k = 1; k < 8; k++) lanes[0] += lanes[k];
for (int i = n - n % 8; i < n; i++) lanes[0] += a[i];

def vector_reduce(data, vf):
    """模拟向量归约: vf 条部分和, 最后横向合并."""
    partial = [0.0] * vf
    for i, x in enumerate(data):
        partial[i % vf] += x
    return sum(partial), partial

print(vector_reduce([1.0] * 100, 8))

// 条件赋值向量化: 用掩码选择
for (int i = 0; i < n; i += 8) {
    __m256 va = _mm256_loadu_ps(a + i);
    __m256 vb = _mm256_loadu_ps(b + i);
    __m256 mask = _mm256_cmp_ps(va, _mm256_setzero_ps(), _CMP_GT_OQ);
    _mm256_storeu_ps(b + i, _mm256_blendv_ps(vb, va, mask));  // 满足则取 a
}

// 手写 intrinsics 版: 点积
#include <immintrin.h>
float dot_avx(const float *a, const float *b, int n) {
    __m256 acc = _mm256_setzero_ps();
    for (int i = 0; i + 8 <= n; i += 8)
        acc = _mm256_fmadd_ps(_mm256_loadu_ps(a + i),
                              _mm256_loadu_ps(b + i), acc);
    float tmp[8];
    _mm256_storeu_ps(tmp, acc);
    float s = tmp[0] + tmp[1] + tmp[2] + tmp[3];
    s += tmp[4] + tmp[5] + tmp[6] + tmp[7];
    for (int i = n - n % 8; i < n; i++) s += a[i] * b[i];
    return s;
}

// 可移植的向量类型: GCC/Clang 的 vector extension
typedef float v8sf __attribute__((vector_size(32)));

v8sf add8(v8sf a, v8sf b) { return a + b; }   // 直接写运算符
float lane0(v8sf v) { return v[0]; }

def speedup_model(vf, mem_bytes_per_iter, compute_cycles,
                  bandwidth_bytes_per_cycle, overhead_ratio=0.15):
    """粗略估算向量化收益: 取访存与计算瓶颈的较大者."""
    compute = compute_cycles / vf * (1 + overhead_ratio)
    memory = mem_bytes_per_iter / bandwidth_bytes_per_cycle
    scalar = compute_cycles + mem_bytes_per_iter / bandwidth_bytes_per_cycle
    return scalar / max(compute, memory)

# 访存密集: VF 提升收益有限
print(round(speedup_model(8, 24, 3, 32), 2))
# 计算密集: 接近 VF
print(round(speedup_model(8, 24, 40, 32), 2))

# 实测: 关掉与打开向量化的对比
cc -O3 -march=native -fno-tree-vectorize bench.c -o bench_scalar
cc -O3 -march=native                        bench.c -o bench_vector
perf stat -e cycles,instructions,fp_arith_inst_retired.256b ./bench_vector