cpu-optimization-x64

x64 CPU 性能优化指南

1. x64 架构特性与优化策略

1.1 架构标识

• 架构: x86_64 (也称为 x64, AMD64)
• 主要厂商: Intel, AMD
• SIMD 扩展: AVX, AVX2, AVX-512

1.2 核心优化原则

1. 利用 SIMD 并行性: 使用 AVX/AVX2/AVX-512 指令同时处理多个数据
2. 优化缓存使用: 按行优先访问，提高缓存命中率
3. 减少分支预测失败: 循环展开，减少条件判断
4. 内存对齐: 确保数据对齐到 32/64 字节边界

2. SIMD/AVX 向量化优化

2.1 基本概念

AVX (Advanced Vector Extensions) 是 x86-64 的 SIMD 指令集扩展：

• AVX: 256 位寄存器，可同时处理 8 个 float32 或 4 个 float64
• AVX2: 增强的 AVX，支持整数运算
• AVX-512: 512 位寄存器，可同时处理 16 个 float32 或 8 个 float64

2.2 编译器自动向量化

推荐方式: 让编译器自动向量化，通过编译选项启用：

# 在 load_inline 中添加向量化选项
op_module = load_inline(
    name="custom_op",
    cpp_sources=cpp_source,
    extra_cflags=[
        "-O3",              # 最高优化级别
        "-march=native",    # 针对当前 CPU 架构优化
        "-ftree-vectorize", # 启用自动向量化
    ],
    verbose=True
)

2.3 循环优化示例

简单方式（未优化）:

torch::Tensor elementwise_add(torch::Tensor a, torch::Tensor b) {
    if (!a.is_contiguous()) a = a.contiguous();
    if (!b.is_contiguous()) b = b.contiguous();
    
    torch::Tensor output = torch::zeros_like(a);
    auto a_ptr = a.data_ptr<float>();
    auto b_ptr = b.data_ptr<float>();
    auto out_ptr = output.data_ptr<float>();
    int64_t numel = a.numel();
    
    // 简单循环
    for (int64_t i = 0; i < numel; ++i) {
        out_ptr[i] = a_ptr[i] + b_ptr[i];
    }
    
    return output;
}

优化方式（循环展开，便于向量化）:

torch::Tensor elementwise_add_optimized(torch::Tensor a, torch::Tensor b) {
    if (!a.is_contiguous()) a = a.contiguous();
    if (!b.is_contiguous()) b = b.contiguous();
    
    torch::Tensor output = torch::zeros_like(a);
    auto a_ptr = a.data_ptr<float>();
    auto b_ptr = b.data_ptr<float>();
    auto out_ptr = output.data_ptr<float>();
    int64_t numel = a.numel();
    
    // 循环展开 8 倍（匹配 AVX 寄存器宽度）
    int64_t i = 0;
    int64_t step = 8;
    for (; i + step <= numel; i += step) {
        out_ptr[i]     = a_ptr[i]     + b_ptr[i];
        out_ptr[i + 1] = a_ptr[i + 1] + b_ptr[i + 1];
        out_ptr[i + 2] = a_ptr[i + 2] + b_ptr[i + 2];
        out_ptr[i + 3] = a_ptr[i + 3] + b_ptr[i + 3];
        out_ptr[i + 4] = a_ptr[i + 4] + b_ptr[i + 4];
        out_ptr[i + 5] = a_ptr[i + 5] + b_ptr[i + 5];
        out_ptr[i + 6] = a_ptr[i + 6] + b_ptr[i + 6];
        out_ptr[i + 7] = a_ptr[i + 7] + b_ptr[i + 7];
    }
    
    // 处理剩余元素
    for (; i < numel; ++i) {
        out_ptr[i] = a_ptr[i] + b_ptr[i];
    }
    
    return output;
}