cpu-optimization-arm

ARM CPU 性能优化指南

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

1.1 架构标识

• 架构: aarch64 (ARM 64-bit, ARMv8-A)
• 主要厂商: ARM, Apple Silicon (M1/M2/M3), AWS Graviton, 华为鲲鹏
• SIMD 扩展: NEON (Advanced SIMD)

1.2 核心优化原则

1. 利用 NEON 并行性: 使用 NEON 指令同时处理多个数据
2. 消除数据依赖: 避免连续指令间的寄存器依赖
3. 优化缓存使用: 按行优先访问，提高缓存命中率
4. 减少分支预测失败: 循环展开，减少条件判断

2. NEON SIMD 向量化优化

2.1 基本概念

NEON (Advanced SIMD) 是 ARM 的 SIMD 指令集：

• 寄存器宽度: 128 位
• 并行处理能力:
  • 4 个 float32（单精度浮点）
  • 2 个 float64（双精度浮点）
  • 16 个 int8, 8 个 int16, 4 个int32, 2 个 int64

2.2 编译器自动向量化

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

# 在 load_inline 中添加 ARM 向量化选项
op_module = load_inline(
    name="custom_op",
    cpp_sources=cpp_source,
    extra_cflags=[
        "-O3",                  # 最高优化级别
        "-mcpu=native",         # 针对当前 ARM CPU 优化
        "-ftree-vectorize",     # 启用自动向量化
        "-ffast-math",          # 快速数学优化（可选）
    ],
    verbose=True
)

注意: ARM 使用 -mcpu=native 而不是 -march=native。

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;
}

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

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();
    
    // 循环展开 4 倍（匹配 NEON 对 float32 的处理能力）
    int64_t i = 0;
    int64_t step = 4;
    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];
    }
    
    // 处理剩余元素
    for (; i < numel; ++i) {
        out_ptr[i] = a_ptr[i] + b_ptr[i];
    }
    
    return output;
}