跳转至

# ARM与NEON

ARM处理器驱动着每一部智能手机、大多数平板电脑、Apple的笔记本电脑以及日益增长的数据中心服务器份额。本文涵盖ARM架构、使用C++内联函数的NEON SIMD编程、用于可伸缩向量处理的SVE/SVE2、Apple Silicon特性以及实际向量化核函数示例

- 如果你拥有iPhone、MacBook或使用AWS Graviton实例，你正在运行ARM。ARM的功耗效率使其在移动和嵌入式领域占据主导地位，并在服务器和ML推理方面日益具有竞争力。理解ARM SIMD让你能够编写在大多数人实际使用的硬件上快速运行的代码。

- 有关生产中ARM SIMD核函数的实际例子，请参见Cactus——面向移动设备和可穿戴设备的低延迟AI引擎：github.com/cactus-compute/cactus。Cactus实现了自定义ARM NEON和NPU加速的注意机制、KV缓存量化和分块预填充核函数，在ARM CPU上实现了最快的推理，且RAM比其它引擎低10倍。其三层架构（引擎→图→核函数）是本文中SIMD概念如何用于构建生产级ML基础设施的具体实例。

## ARM架构基础

- ARM是一种RISC（精简指令集计算机）架构（第13章）。关键特征：

  - 加载-存储架构：算术指令只操作寄存器，从不直接操作内存。要对内存中的两个数相加，你必须：(1) 将它们加载到寄存器，(2) 将寄存器相加，(3) 将结果存回内存。这比x86更简单（x86可以在一条指令中加一个寄存器和一个内存位置），但使得流水线更清晰。

  - 定长指令：每个ARMv8（AArch64）指令恰好32位。这使得解码快速且可预测（不像x86的可变长指令，长度可以是1-15字节）。

  - 32个通用寄存器（x0-x30，每个64位）加上栈指针（sp）和零寄存器（xzr）。相比之下x86有16个通用寄存器。更多寄存器 = 更少内存访问 = 更快代码。

  - 32个SIMD/浮点寄存器（v0-v31，每个128位）用于NEON和浮点操作。

    // ARM汇编（仅感受风格——你将使用内联函数，而非汇编）
    // 两寄存器相加
    add x0, x1, x2    // x0 = x1 + x2

    // 从内存加载
    ldr x0, [x1]      // x0 = *x1（从x1中的地址加载64位）

    // NEON：加四个浮点数
    fadd v0.4s, v1.4s, v2.4s  // v0 = v1 + v2（四个32位浮点数）

- 你不会写汇编。你将使用内联函数：与特定指令一对一映射的C/C++函数。编译器处理寄存器分配、调度和其他底层细节。

## NEON：128位SIMD

- NEON是ARM的SIMD扩展。每个NEON寄存器宽128位，可容纳：

数据类型 | 每寄存器元素数 | 表示法
--- | --- | ---
float32 | 4 | `float32x4_t`
float16 | 8 | `float16x8_t`
int32 | 4 | `int32x4_t`
int16 | 8 | `int16x8_t`
int8 | 16 | `int8x16_t`

- 128位比x86的AVX（256位）或AVX-512（512位）窄。但ARM以出色的功耗效率和广泛的可用性弥补了这一点。

### NEON内联函数：基础

- NEON内联函数遵循命名约定：`v[操作][限定符]_[类型]`

    #include <arm_neon.h>

    // 从内存加载4个浮点数到NEON寄存器
    float32x4_t a = vld1q_f32(ptr);        // vld1q = vector load 1, q = 128位（四字）

    // 从NEON寄存器存储4个浮点数到内存
    vst1q_f32(out_ptr, a);                   // vst1q = vector store 1, q = 128位

    // 算术运算
    float32x4_t c = vaddq_f32(a, b);        // c = a + b（4个浮点数）
    float32x4_t d = vmulq_f32(a, b);        // d = a * b（4个浮点数）
    float32x4_t e = vfmaq_f32(c, a, b);     // e = c + a * b（融合乘加，4个浮点数）

    // 比较（返回掩码：若真则全1，若假则全0）
    uint32x4_t mask = vcgtq_f32(a, b);      // mask[i] = (a[i] > b[i]) ? 0xFFFFFFFF : 0

    // 基于掩码选择元素（类似numpy.where）
    float32x4_t result = vbslq_f32(mask, a, b);  // result[i] = mask[i] ? a[i] : b[i]

    // 归约：将所有4个元素求和为标量
    float total = vaddvq_f32(a);             // total = a[0] + a[1] + a[2] + a[3]

- `vfmaq_f32`（融合乘加）是ML最重要的SIMD指令。它用一次舍入步骤计算 \(c = c + a \times b\)（比分开乘然后加更精确）。点积、矩阵乘法和卷积都由FMA构建。

### 实践示例：向量化点积

- 点积是矩阵乘法的内循环。让我们先用标量C++编写，然后用NEON向量化。

    #include <arm_neon.h>

    // 标量点积
    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;
    }

    // NEON向量化点积
    float dot_neon(const float* a, const float* b, int n) {
        float32x4_t sum_vec = vdupq_n_f32(0.0f);  // 初始化4个累加器为0

        int i = 0;
        for (; i + 4 <= n; i += 4) {
            float32x4_t va = vld1q_f32(a + i);     // 从a加载4个元素
            float32x4_t vb = vld1q_f32(b + i);     // 从b加载4个元素
            sum_vec = vfmaq_f32(sum_vec, va, vb);   // sum_vec += va * vb
        }

        // 将4个累加器归约为单一标量
        float sum = vaddvq_f32(sum_vec);

        // 处理剩余元素（如果n不是4的倍数）
        for (; i < n; i++) {
            sum += a[i] * b[i];
        }

        return sum;
    }

- 关键C++概念：

  - `const float*`：指向只读浮点数据的指针。`const` 承诺我们不会通过此指针修改数据。
  - `a + i`：指针运算。`a + i` 指向数组的第 \(i\) 个元素（等同于 `&a[i]`）。
  - 末尾的"清理循环"处理 \(n\) 不是4的倍数的情况。这是SIMD代码中的通用模式：用向量化块处理主体部分，然后用标量代码处理余数。

- 为什么 `sum_vec` 中使用4个累加器：我们使用4个独立的累加器（每个SIMD通道一个），而不是单个标量累加器。这避免了数据依赖：每次迭代的FMA依赖于 `sum_vec`，但有了4个独立通道，CPU可以对FMAs进行流水线处理。最后，我们将4个部分和归约为一个。

### 实践示例：向量化ReLU

    #include <arm_neon.h>

    void relu_neon(const float* input, float* output, int n) {
        float32x4_t zero = vdupq_n_f32(0.0f);

        int i = 0;
        for (; i + 4 <= n; i += 4) {
            float32x4_t x = vld1q_f32(input + i);
            float32x4_t result = vmaxq_f32(x, zero);  // max(x, 0) = ReLU
            vst1q_f32(output + i, result);
        }

        // 标量清理
        for (; i < n; i++) {
            output[i] = input[i] > 0 ? input[i] : 0;
        }
    }

- `vmaxq_f32` 计算两个向量的逐元素最大值。由于一个向量全为零，这恰好就是ReLU。无需分支，无需比较——仅一条指令。

## I8MM：整数矩阵乘法

- I8MM（Int8矩阵乘法）是ARMv8.6扩展，增加了用于INT8矩阵乘法（INT32累加）的专用指令——这正是量化ML推理所需要的。

- 关键指令是 `SMMLA`（有符号矩阵乘加）：它接受两个8×2块的INT8值，并将结果累加到2×2块的INT32中：

    #include <arm_neon.h>

    // I8MM：将两个8元素INT8向量相乘，累加到4个INT32结果中
    // 这从2x8 × 8x2输入块计算输出矩阵的一个2x2瓦片
    void matmul_i8mm_tile(const int8_t* A, const int8_t* B, int32_t* C) {
        // 从A加载8字节（2行各4元素，打包）
        int8x16_t va = vld1q_s8(A);   // 16字节 = 2行 × 8元素
        int8x16_t vb = vld1q_s8(B);   // 16字节 = 2行 × 8元素

        // 加载现有累加器（2x2 = 4个int32值）
        int32x4_t acc = vld1q_s32(C);

        // I8MM指令：acc += A_tile × B_tile^T
        // 从2×8 × 8×2输入计算2×2输出
        acc = vmmlaq_s32(acc, va, vb);  // I8MM指令

        vst1q_s32(C, acc);
    }

- 为什么I8MM重要：没有I8MM时，NEON上的INT8矩阵乘法需要加宽乘法（`vmull`）后跟成对加法——每个输出元素需要多条指令。有了I8MM，硬件在一条指令中完成8元素点积（2×8 × 8×2 = 2×2）。对于INT8推理工作负载，这比纯NEON快4-8倍。

- 可用性：Apple M1+（所有Apple Silicon）、ARM Cortex-A510/A710/X2+（ARMv9）、AWS Graviton3+。用 `#ifdef __ARM_FEATURE_MATMUL_INT8` 检查。

- 对于ML推理：在ARM服务器（Graviton）或Apple Silicon上运行的INT8量化模型（第18章）从I8MM中获益巨大。ONNX Runtime和llama.cpp等框架在运行时检测I8MM并自动使用优化核函数。

## SME和SME2：可伸缩矩阵扩展

- SME（可伸缩矩阵扩展）是ARM对Intel AMX和NVIDIA张量核心的回应：用于矩阵操作的专用硬件。SME2（ARMv9.2）进一步扩展了它。

- SME引入了ZA瓦片寄存器：存储在硬件中的2D矩阵，最大可达SVL×SVL字节（其中SVL是流向量长度，通常每维128-512位）。与NEON（1D向量）甚至SVE（1D可伸缩向量）不同，SME原生操作2D瓦片。

- 编程模型有两种模式：

  - 普通模式：标准ARM执行（NEON、SVE正常工作）。
  - 流SVE模式：通过 `smstart` 进入，启用SME指令。SVE指令在此模式下也可工作，但可能使用不同的寄存器宽度。

    #include <arm_sme.h>

    // SME2：矩阵乘法的外积累加
    // 将A_col × B_row 累加到ZA瓦片寄存器中
    void sme2_matmul_outer(const float* A_col, const float* B_row, int K) {
        // 进入流模式
        // smstart;  // （通过编译器内联或内联汇编完成）

        // 清零ZA瓦片累加器
        svzero_za();

        for (int k = 0; k < K; k++) {
            // 将A的一列和B的一行加载到SVE寄存器中
            svfloat32_t a = svld1_f32(svptrue_b32(), &A_col[k * SVL]);
            svfloat32_t b = svld1_f32(svptrue_b32(), &B_row[k * SVL]);

            // 外积：ZA += a × b^T
            // 这在一个指令中累加一个SVL×SVL瓦片
            svmopa_za32_f32_m(0, svptrue_b32(), svptrue_b32(), a, b);
        }

        // 将ZA瓦片存储到内存
        // svst1_za(...);

        // 退出流模式
        // smstop;
    }

- 关键概念：

  - `svmopa`（外积累加）：核心SME指令。它计算两个向量的完整外积并累加到ZA瓦片中。对于SVL=512位（16个浮点数），这是一个16×16外积——一条指令中256次FMA操作。
  - ZA瓦片：在流模式中跨指令持久存在。你将多个外积（每个K迭代一个）累加到同一瓦片中，构建完整的矩阵乘法瓦片。
  - 流模式：SME指令仅在流模式下工作。进入/退出流模式的开销意味着SME最适合持续的矩阵计算，而非短时爆发。

- SME2新增：多向量操作（同时处理2或4个SVE向量）、额外的瓦片操作以及与普通模式的改进集成。

- 可用性：ARM Neoverse V2（AWS Graviton4）、一些即将推出的移动芯片。截至2026年尚未出现在Apple Silicon上。SME仍处于早期阶段——大多数ML框架还没有SME优化的核函数。

- 演进脉络：NEON（128位向量，逐元素）→ I8MM（INT8矩阵瓦片）→ SVE（可伸缩向量）→ SME（可伸缩2D矩阵瓦片）。每一代都更接近硬件原生矩阵操作。

## SVE和SVE2：可伸缩向量扩展

- NEON具有固定的128位宽度。SVE（可伸缩向量扩展）引入了向量长度无关（VLA）编程：你编写一次代码，它在任何向量宽度（128到2048位）的硬件上运行。硬件在运行时确定宽度。

    #include <arm_sve.h>

    void add_sve(const float* a, const float* b, float* c, int n) {
        int i = 0;
        svbool_t pred = svwhilelt_b32(i, n);  // 谓词：哪些通道是激活的

        while (svptest_any(svptrue_b32(), pred)) {
            svfloat32_t va = svld1(pred, a + i);
            svfloat32_t vb = svld1(pred, b + i);
            svst1(pred, c + i, svadd_x(pred, va, vb));

            i += svcntw();  // 按硬件向量宽度前进（以32位元素计）
            pred = svwhilelt_b32(i, n);
        }
    }

- 谓词寄存器（`svbool_t`）取代了标量清理循环。每个通道有一个谓词位：激活的通道参与，非激活的被屏蔽。`svwhilelt_b32(i, n)` 指令创建一个谓词，其中对应 `i, i+1, ..., n-1` 的通道被激活。这自动处理了尾部。

- `svcntw()` 在运行时返回每个向量寄存器中32位元素的数量。在具有256位SVE的CPU上，返回8。在512位SVE上，返回16。你的代码自动适应。

- SVE在ARM Neoverse V1/V2上可用（AWS Graviton3/4，一些服务器芯片）。在Apple Silicon上尚不可用。

## Apple Silicon特性

- Apple的M系列芯片（M1、M2、M3、M4）是基于ARM的自定义微架构：

- 性能核心和效率核心：P核心（Firestorm/Avalanche等）用于重型计算，E核心（Icestorm/Blizzard等）用于后台任务。调度器将线程分配给适当的核心类型。

- AMX（Apple矩阵扩展）：专用矩阵乘法单元，独立于NEON。AMX未公开（Apple不发布ISA），但Accelerate框架内部将其用于BLAS操作。当你在Mac上调用 `np.dot` 时，它通过Accelerate，后者使用AMX。你不能直接对AMX编程（除非逆向工程）。

- 统一内存：CPU和GPU共享同一物理RAM。在其他系统上，数据必须从CPU内存拷贝到GPU内存（通过PCIe，约32 GB/s）。在Apple Silicon上，无需拷贝——GPU读取CPU写入的同一内存。这消除了ML工作负载的主要瓶颈。

- 神经网络引擎：一个16核专用ML加速器。INT8推理时达到约30 TOPS（每秒万亿次操作）。Core ML将其用于设备端推理。

- Apple Silicon上的ML：使用MLX（Apple的ML框架），它专为统一内存架构设计。PyTorch也有MPS（Metal性能着色器）后端支持，尽管不如CUDA成熟。

## 自动向量化

- 编写SIMD内联函数很繁琐。编译器能自动向量化你的代码吗？

- 可以的，但有限制。现代编译器（GCC、Clang）可以自动向量化简单循环：

    // 编译器可以自动向量化此代码（使用 -O3 -march=native）
    void add_auto(const float* a, const float* b, float* c, int n) {
        for (int i = 0; i < n; i++) {
            c[i] = a[i] + b[i];
        }
    }

- 有助于自动向量化的模式：
  - 简单的循环，迭代次数已知。
  - 迭代之间无数据依赖（`c[i]` 不依赖于 `c[i-1]`）。
  - 连续内存访问（无分散/聚集）。
  - `const` 和 `restrict` 指针（告知编译器数组不重叠）。

    // restrict 告诉编译器：a、b、c 指向不重叠的内存
    void add_restrict(const float* __restrict__ a,
                      const float* __restrict__ b,
                      float* __restrict__ c, int n) {
        for (int i = 0; i < n; i++) {
            c[i] = a[i] + b[i];
        }
    }

- 没有 `restrict`，编译器必须假设 `c` 可能与 `a` 或 `b` 重叠（写入 `c[i]` 可能改变 `a[i+1]`），从而阻止向量化。

- 阻止自动向量化的模式：

  - 数据依赖：`a[i] = a[i-1] + b[i]`（每次迭代依赖前一次）。
  - 复杂控制流：循环内的 `if` 语句（除非编译器能转换为谓词化）。
  - 循环内的函数调用（除非函数被内联）。
  - 指针别名（数组可能重叠，没有 `restrict`）。

- 检查自动向量化：使用编译器标志查看哪些被向量化了：

    # GCC：显示向量化决策
    g++ -O3 -march=native -fopt-info-vec-optimized code.cpp

    # Clang：显示向量化报告
    clang++ -O3 -march=native -Rpass=loop-vectorize code.cpp

- 何时使用内联函数 vs 自动向量化：从干净的C++和编译器优化开始。如果编译器向量化了你的循环，很好。如果性能仍不足，检查编译器的向量化报告以理解原因，然后仅为关键内循环编写内联函数。过早使用内联函数会让代码难以阅读而没有确定的收益。

## 编程任务（在ARM上用g++或clang++编译——Mac M系列或Linux aarch64）

1. 编写标量点积和NEON向量化点积。对两者进行基准测试并测量加速比。

       // task1_neon_dot.cpp
       // 编译（Mac/ARM Linux）：clang++ -O3 -o task1 task1_neon_dot.cpp
       // 注意：NEON在AArch64上默认启用，无需特殊标志

       #include <iostream>
       #include <chrono>
       #include <vector>
       #include <arm_neon.h>

       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_neon(const float* a, const float* b, int n) {
           float32x4_t sum_vec = vdupq_n_f32(0.0f);
           int i = 0;
           for (; i + 4 <= n; i += 4) {
               float32x4_t va = vld1q_f32(a + i);
               float32x4_t vb = vld1q_f32(b + i);
               sum_vec = vfmaq_f32(sum_vec, va, vb);
           }
           float sum = vaddvq_f32(sum_vec);
           for (; i < n; i++) sum += a[i] * b[i];
           return sum;
       }

       int main() {
           const int N = 10'000'000;
           std::vector<float> a(N, 1.0f), b(N, 2.0f);

           // 预热
           volatile float s1 = dot_scalar(a.data(), b.data(), N);
           volatile float s2 = dot_neon(a.data(), b.data(), N);

           // 标量基准测试
           auto start = std::chrono::high_resolution_clock::now();
           for (int t = 0; t < 100; t++) {
               s1 = dot_scalar(a.data(), b.data(), N);
           }
           auto end = std::chrono::high_resolution_clock::now();
           double scalar_ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;

           // NEON基准测试
           start = std::chrono::high_resolution_clock::now();
           for (int t = 0; t < 100; t++) {
               s2 = dot_neon(a.data(), b.data(), N);
           }
           end = std::chrono::high_resolution_clock::now();
           double neon_ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;

           std::cout << "标量: " << scalar_ms << " ms（结果: " << s1 << "）\n";
           std::cout << "NEON: " << neon_ms << " ms（结果: " << s2 << "）\n";
           std::cout << "加速比: " << scalar_ms / neon_ms << "x\n";
           return 0;
       }

2. 实现NEON ReLU和softmax最大值查找。练习使用不同操作的加载→计算→存储模式。

       // task2_neon_ops.cpp
       // 编译：clang++ -O3 -o task2 task2_neon_ops.cpp

       #include <iostream>
       #include <vector>
       #include <cmath>
       #include <arm_neon.h>

       void relu_neon(const float* in, float* out, int n) {
           float32x4_t zero = vdupq_n_f32(0.0f);
           int i = 0;
           for (; i + 4 <= n; i += 4) {
               float32x4_t x = vld1q_f32(in + i);
               vst1q_f32(out + i, vmaxq_f32(x, zero));
           }
           for (; i < n; i++) out[i] = in[i] > 0 ? in[i] : 0;
       }

       float max_neon(const float* data, int n) {
           float32x4_t max_vec = vdupq_n_f32(-INFINITY);
           int i = 0;
           for (; i + 4 <= n; i += 4) {
               max_vec = vmaxq_f32(max_vec, vld1q_f32(data + i));
           }
           float result = vmaxvq_f32(max_vec);
           for (; i < n; i++) result = result > data[i] ? result : data[i];
           return result;
       }

       int main() {
           std::vector<float> data = {-3, 1, -1, 4, 2, -5, 0, 7, -2, 3};
           std::vector<float> out(data.size());

           relu_neon(data.data(), out.data(), data.size());
           std::cout << "ReLU: ";
           for (float x : out) std::cout << x << " ";
           std::cout << "\n";

           float mx = max_neon(data.data(), data.size());
           std::cout << "最大值: " << mx << "（期望值: 7）\n";
           return 0;
       }

3. 比较自动向量化代码与手写NEON内联函数。用 `-fopt-info-vec`（GCC）或 `-Rpass=loop-vectorize`（Clang）编译以查看编译器的操作。

       // task3_auto_vs_manual.cpp
       // 编译：clang++ -O3 -Rpass=loop-vectorize -o task3 task3_auto_vs_manual.cpp
       //    （或）：g++ -O3 -fopt-info-vec-optimized -o task3 task3_auto_vs_manual.cpp

       #include <iostream>
       #include <chrono>
       #include <vector>
       #include <arm_neon.h>

       // 让编译器自动向量化
       void add_auto(const float* __restrict__ a, const float* __restrict__ b,
                     float* __restrict__ c, int n) {
           for (int i = 0; i < n; i++) {
               c[i] = a[i] + b[i];
           }
       }

       // 手写NEON
       void add_neon(const float* a, const float* b, float* c, int n) {
           int i = 0;
           for (; i + 4 <= n; i += 4) {
               vst1q_f32(c + i, vaddq_f32(vld1q_f32(a + i), vld1q_f32(b + i)));
           }
           for (; i < n; i++) c[i] = a[i] + b[i];
       }

       int main() {
           const int N = 10'000'000;
           std::vector<float> a(N, 1.0f), b(N, 2.0f), c(N);

           auto bench = [&](auto fn, const char* name) {
               fn(a.data(), b.data(), c.data(), N);  // 预热
               auto start = std::chrono::high_resolution_clock::now();
               for (int t = 0; t < 100; t++) fn(a.data(), b.data(), c.data(), N);
               auto end = std::chrono::high_resolution_clock::now();
               double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
               std::cout << name << ": " << ms << " ms\n";
           };

           bench(add_auto, "自动向量化");
           bench(add_neon, "手写NEON");
           // 它们应该非常接近——编译器能很好地自动向量化这个简单循环
           return 0;
       }