ARM 与 NEON

文档摘要

ARM 与 NEON ARM 处理器驱动着每一部智能手机、大部分平板、苹果的笔记本，以及数据中心服务器中越来越大的份额。本节讲 ARM 架构、用 C++ 内联函数做 NEON SIMD 编程、用于可扩展向量处理的 SVE/SVE2、Apple Silicon 的特性，以及实用的向量化内核示例。只要你手上有 iPhone、MacBook，或在用 AWS Graviton 实例，你跑的就是 ARM。ARM 的能效让它在移动和嵌入式领域占据主导，在服务器和 ML 推理中也愈发有竞争力。理解 ARM SIMD，你就能为大多数人实际使用的硬件写出跑得快的代码。想看一个生产环境中的 ARM SIMD 内核实例，可以看看 Cactus——一个面向移动设备和可穿戴设备的低延迟 AI 引擎：github.

ARM 与 NEON

ARM 处理器驱动着每一部智能手机、大部分平板、苹果的笔记本，以及数据中心服务器中越来越大的份额。本节讲 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 为注意力机制、KV-cache 量化、分块预填充实现了自定义的 ARM NEON 与 NPU 加速内核，在 ARM CPU 上做到了最快推理，且内存占用比其他引擎低 10 倍。它的三层架构（Engine → Graph → Kernels）是本节 SIMD 概念如何用于构建生产级 ML 基础设施的一个具体范例。

ARM 架构基础

ARM 是一种 RISC（Reduced Instruction Set Computer，精简指令集计算机）架构（第 13 章）。关键特征：
- 加载-存储架构（load-store architecture）：算术指令只操作寄存器，绝不直接操作内存。要把内存里的两个数相加，你必须：(1) 把它们加载到寄存器，(2) 把寄存器相加，(3) 把结果存回内存。这比 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：四个 float 相加
fadd v0.4s, v1.4s, v2.4s  // v0 = v1 + v2（四个 32 位 float）

你不会去写汇编。你会用内联函数（intrinsics）：与具体指令一一对应的 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 个 float 到一个 NEON 寄存器
float32x4_t a = vld1q_f32(ptr);        // vld1q = 向量 load 1，q = 128 位（quad）

// 把 4 个 float 从 NEON 寄存器存到内存
vst1q_f32(out_ptr, a);                   // vst1q = 向量 store 1，q = 128 位

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

// 比较（返回掩码：真则全 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（融合乘加，fused multiply-add）是 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*：指向只读 float 数据的指针。const 承诺我们不会通过这个指针修改数据。
- a + i：指针算术。a + i 指向数组第 i 个元素（等价于 &a[i]）。
- 末尾的「收尾循环」处理 n 不是 4 的倍数的情况。这是 SIMD 代码中的通用模式：用向量化块处理主体，再用标量代码处理余数。
为什么 sum_vec 里有 4 个累加器：我们不用单个标量累加器，而是用 4 个相互独立的累加器（每条 SIMD 通道一个）。这避免了数据依赖：每次迭代的 FMA 都依赖 sum_vec，但因为有 4 条独立通道，CPU 可以把 FMA 流水化。最后再把 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 计算两个向量的逐元素最大值。由于其中一个向量全为 0，这正好就是 ReLU。没有分支、没有比较——只有一条指令。

I8MM：整数矩阵乘法

I8MM（Int8 Matrix Multiply，Int8 矩阵乘）是 ARMv8.6 的扩展，它加入了带 INT32 累加的 INT8 矩阵乘专用指令——恰好是量化 ML 推理所需的东西。
关键指令是 SMMLA（Signed Matrix Multiply-Accumulate，有符号矩阵乘累加）：它取两块 8×2 的 INT8 值，把结果累加进一块 2×2 的 INT32：


#include <arm_neon.h>

// I8MM：两个 8 元素 INT8 向量相乘，累加到 4 个 INT32 结果
// 它从 2x8 x 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 matmul 需要先做宽位乘法（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 推理：INT8 量化模型（第 18 章）跑在 ARM 服务器（Graviton）或 Apple Silicon 上时，能从 I8MM 中获益巨大。ONNX Runtime 和 llama.cpp 这类框架会在运行时检测 I8MM 并自动用上优化内核。

SME 与 SME2：可扩展矩阵扩展

SME（Scalable Matrix Extension，可扩展矩阵扩展）是 ARM 对 Intel AMX 和 NVIDIA Tensor Core 的回应：为矩阵运算而生的专用硬件。SME2（ARMv9.2）对它做了进一步扩展。
SME 引入了 ZA tile 寄存器：存在硬件里的二维矩阵，最大可达 SVL×SVL 字节（其中 SVL 是流式向量长度，通常每维 128-512 位）。与 NEON（一维向量）甚至 SVE（一维可扩展向量）不同，SME 原生地操作二维瓦片（tile）。
它的编程模型有两种模式：
- 普通模式（normal mode）：标准 ARM 执行（NEON、SVE 照常工作）。
- 流式 SVE 模式（streaming SVE mode）：通过 smstart 进入，启用 SME 指令。SVE 指令在这种模式下也能用，但可能使用不同的寄存器宽度。


#include <arm_sme.h>

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

    // 把 ZA tile 累加器清零
    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 tile 存回内存
    // svst1_za(...);

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

关键概念：
- svmopa（外积累加，outer product accumulate）：SME 的核心指令。它计算两个向量的完整外积并累加进 ZA tile。当 SVL=512 位（16 个 float）时，这是一个 16×16 的外积——一条指令就是 256 次 FMA。
- ZA tile：在流式模式内，跨指令持久存在。你把多个外积（每次 K 迭代一次）累加进同一个 tile，逐步拼出一个完整的矩阵乘瓦片。
- 流式模式：SME 指令只在流式模式下有效。进出流式模式有开销，所以 SME 最适合持续的矩阵运算，而不是短促的零散使用。
SME2 的增补：多向量操作（一次处理 2 个或 4 个 SVE 向量）、额外的 tile 操作、以及与普通模式更好的集成。
可用性：ARM Neoverse V2（AWS Graviton4）、一些即将上市的移动芯片。截至 2026 年尚未出现在 Apple Silicon 上。SME 仍处于早期阶段——大多数 ML 框架还没有 SME 优化的内核。
演进路线：NEON（128 位向量，逐元素）→ I8MM（INT8 矩阵瓦片）→ SVE（可扩展向量）→ SME（可扩展二维矩阵瓦片）。每一代都更接近「硬件原生地做矩阵运算」。

SVE 与 SVE2：可扩展向量扩展

NEON 的宽度固定为 128 位。SVE（Scalable Vector Extension，可扩展向量扩展）引入了**向量长度无关（vector-length agnostic，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);
    }
}

**谓词寄存器（predicate register，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 特性

苹果的 M 系列芯片（M1、M2、M3、M4）基于 ARM，但带有自定义的微架构：
性能核与能效核：P 核（Firestorm/Avalanche 等）负责重计算，E 核（Icestorm/Blizzard 等）负责后台任务。调度器会把线程分配给合适的核心类型。
AMX（Apple Matrix eXtensions，苹果矩阵扩展）：专用的矩阵乘单元，独立于 NEON。AMX 没有公开文档（苹果不发布它的 ISA），但 Accelerate 框架内部用它做 BLAS。当你在 Mac 上调用 np.dot 时，它会经过 Accelerate，进而用到 AMX。你无法直接编程 AMX（除非反向工程）。
统一内存：CPU 和 GPU 共享同一块物理 RAM。在其他系统上，数据必须从 CPU 内存复制到 GPU 内存（经 PCIe，约 32 GB/s）。在 Apple Silicon 上，没有复制——GPU 直接读 CPU 写过的同一块内存。这消除了 ML 工作负载中的一大瓶颈。
Neural Engine：一个 16 核的专用 ML 加速器。对 INT8 推理可达到约 30 TOPS（每秒万亿次操作）。Core ML 用它做端侧推理。
在 Apple Silicon 上做 ML：用 MLX（苹果的 ML 框架），它是为统一内存架构设计的。PyTorch 也有 MPS（Metal Performance Shaders）后端支持，不过它远不如 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]）。
- 连续内存访问（没有 scatter/gather）。
- 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

何时用内联函数、何时用自动向量化：从干净的 C++ 加编译器优化开始。如果编译器把你的循环向量化了，很好。如果性能还不够，就去查编译器的向量化报告，搞清楚原因，然后再为关键的内层循环写内联函数。过早地写内联函数会让代码不可读，还未必有收益。

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

写一个标量点积和一个 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: " << scalar_ms << " ms (result: " << s1 << ")\n";
    std::cout << "NEON:   " << neon_ms << " ms (result: " << s2 << ")\n";
    std::cout << "Speedup: " << scalar_ms / neon_ms << "x\n";
    return 0;
}

实现 NEON 的 ReLU 和 softmax 的求最大值。在不同操作上练习 load→compute→store 模式。


// 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 << "Max: " << mx << " (expected: 7)\n";
    return 0;
}

把自动向量化的代码和手写 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, "Auto-vectorised");
    bench(add_neon, "Hand-written NEON");
    // 两者应该非常接近——编译器对这个简单循环的自动向量化做得很好
    return 0;
}