ARM和NEON¶

ARM处理器驱动每部智能手机、大多数平板、Apple的笔记本以及越来越多的数据中心服务器。本文件涵盖ARM架构、NEON SIMD编程（C++内联函数）、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使用仅为其他引擎的1/10。其三层架构（Engine → Graph → Kernels）是本文件中SIMD概念如何用于构建生产ML基础设施的具体实例。

ARM架构基础¶

ARM是RISC（精简指令集计算机）架构（第13章）。关键特征：
- 加载-存储架构：算术指令仅操作寄存器，从不直接操作内存。要对来自内存的两个数相加，必须：(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：加四个浮点数
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[operation][qualifier]_[type]

#include <arm_neon.h>

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

// 从NEON寄存器存储4个浮点数到内存
vst1q_f32(out_ptr, a);                   // vst1q = 向量存储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个浮点数）

// 比较（返回掩码：true时全1，false时全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*：指向只读float数据的指针。const承诺不会通过此指针修改数据。
- a + i：指针算术。a + i指向数组的第\(i\)个元素（等价于&a[i]）。
- 末尾的"清理循环"处理\(n\)不是4的倍数的情况。这是SIMD代码中的通用模式：以向量化块处理大批量部分，然后以标量代码处理余数。
为什么sum_vec中使用4个累加器：不是使用单个标量累加器，而是使用4个独立累加器（每个SIMD通道一个）。这避免了数据依赖：每次迭代的FMA依赖于sum_vec，但使用4个独立通道，CPU可以对FMA进行流水线处理。最后，将4个部分和归约为1个。

实践示例：向量化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 x 8x2的输入块计算输出矩阵的一个2x2块
void matmul_i8mm_tile(const int8_t* A, const int8_t* B, int32_t* C) {
    // 从A加载8字节（每行4个元素的2行，打包）
    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 Tensor Cores的回应：用于矩阵运算的专用硬件。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，而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（Apple的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

何时使用内联函数 vs 自动向量化：从干净的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_ms << " ms (结果: " << s1 << ")\n";
    std::cout << "NEON:   " << neon_ms << " ms (结果: " << s2 << ")\n";
    std::cout << "加速比: " << scalar_ms / neon_ms << "x\n";
    return 0;
}

实现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 << "Max: " << mx << " (期望: 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, "自动向量化");
    bench(add_neon, "手写NEON");
    // 应该非常接近——编译器很好地自动向量化了这个简单循环
    return 0;
}