maths-cs-ai-compendium-zh/chapter 16: SIMD and GPU programming/02. ARM and NEON.md

# 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](https://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和浮点操作。

```cpp
// 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[操作][限定符]_[类型]`

```cpp
#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向量化。

```cpp
#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

```cpp
#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中：

```cpp
#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指令在此模式下也可工作，但可能使用不同的寄存器宽度。

```cpp
#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位）的硬件上运行。硬件在运行时确定宽度。

```cpp
#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）可以自动向量化简单循环：

```cpp
// 编译器可以自动向量化此代码（使用 -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` 指针（告知编译器数组不重叠）。

```cpp
// 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`）。

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

```bash
# 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向量化点积。对两者进行基准测试并测量加速比。
```cpp
// 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最大值查找。练习使用不同操作的加载→计算→存储模式。
```cpp
// 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）编译以查看编译器的操作。
```cpp
// 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;
}
```