ARM 和 NEON¶

ARM 处理器为每部智能手机、大多数平板电脑、Apple 笔记本电脑以及越来越多的数据中心服务器提供支持。该文件涵盖了 ARM 架构、使用 C++ intrinsics、SVE/SVE2 进行可扩展 vector 处理的 NEON SIMD 编程、Apple Silicon 细节以及实用的 vectorised kernel 示例

如果您拥有 iPhone、MacBook 或使用 AWS Graviton 实例，则您正在运行 ARM。 ARM的功效使其在移动和嵌入式领域占据主导地位，并且在服务器和ML inference领域的竞争力日益增强。了解 ARM SIMD 可以让您编写在大多数人实际使用的硬件上快速运行的代码。
有关生产中的 ARM SIMD kernels 的真实示例，请参阅 Cactus - 用于移动设备和可穿戴设备的低延迟 AI 引擎：github.com/cactus-compute/cactus。 Cactus 为 attention、KV-cache 量化和分block预填充实现了自定义 ARM NEON 和 NPU 加速 kernels，在 ARM CPU 上实现了最快的 inference，其 RAM 比其他引擎低 10 倍。其三层架构（引擎 → 图形 → kernel）是如何使用此文件中的 SIMD 概念构建生产 ML 基础设施的具体示例。

ARM 架构基础知识¶

ARM 是一个 RISC（精简指令集计算机）架构（第 13 章）。主要特点：
- 加载-存储架构：算术指令仅在 registers 上运行，从不直接在内存上运行。要从内存中添加两个数字，您必须：(1) 将它们加载到 registers，(2) 将 registers 相加，(3) 将结果存储回内存。这比 x86（可以在一条指令中添加 register 和内存位置）更简单，但可以实现更清晰的流水线操作。
- 固定宽度指令：每个 ARMv8 (AArch64) 指令恰好是 32 位。这使得解码快速且可预测（与 x86 的可变长度指令可以是 1-15 个字节不同）。
- 32个通用registers（x0-x30，每个64位）加上堆栈指针（sp）和零register（xzr）。与x86的16位通用registers相比。更多 registers = 更少的内存访问 = 更快的代码。
- 32 SIMD/floating-point registers（v0-v31，每个 128 位）用于 NEON 和 floating-point 操作。

// ARM assembly (just to see the flavour -- you will use intrinsics, not assembly)
// Add two registers
add x0, x1, x2    // x0 = x1 + x2

// Load from memory
ldr x0, [x1]      // x0 = *x1 (load 64 bits from address in x1)

// NEON: add four floats
fadd v0.4s, v1.4s, v2.4s  // v0 = v1 + v2 (four 32-bit floats)

你不会写汇编。您将使用 intrinsics: C/C++ 函数将 1:1 映射到特定指令。编译器处理 register 分配、调度和其他低级细节。

NEON：128位SIMD¶

NEON是ARM的SIMD extension。每个 NEON register 都是 128 位宽，可以容纳：

数据类型	每个 register 的元素	符号
浮动32	4	`float32x4_t`
浮动16	8	`float16x8_t`
整型32	4	`int32x4_t`
整型16	8	`int16x8_t`
整型8	16	`int8x16_t`

128 位比 x86 的 AVX（256 位）或 AVX-512（512 位）窄。但 ARM 凭借出色的功效和广泛的可用性进行了补偿。

NEON 内在函数：基础知识¶

NEON intrinsics 遵循命名约定：v[operation][qualifier]_[type]

#include <arm_neon.h>

// Load 4 floats from memory into a NEON register
float32x4_t a = vld1q_f32(ptr);        // vld1q = vector load 1, q = 128-bit (quad)

// Store 4 floats from a NEON register to memory
vst1q_f32(out_ptr, a);                   // vst1q = vector store 1, q = 128-bit

// Arithmetic
float32x4_t c = vaddq_f32(a, b);        // c = a + b (4 floats)
float32x4_t d = vmulq_f32(a, b);        // d = a * b (4 floats)
float32x4_t e = vfmaq_f32(c, a, b);     // e = c + a * b (fused multiply-add, 4 floats)

// Comparison (returns a mask: all 1s if true, all 0s if false)
uint32x4_t mask = vcgtq_f32(a, b);      // mask[i] = (a[i] > b[i]) ? 0xFFFFFFFF : 0

// Select elements based on mask (like numpy.where)
float32x4_t result = vbslq_f32(mask, a, b);  // result[i] = mask[i] ? a[i] : b[i]

// Reduce: sum all 4 elements to a scalar
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>

// Scalar dot product
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-vectorised dot product
float dot_neon(const float* a, const float* b, int n) {
    float32x4_t sum_vec = vdupq_n_f32(0.0f);  // initialise 4 accumulators to 0

    int i = 0;
    for (; i + 4 <= n; i += 4) {
        float32x4_t va = vld1q_f32(a + i);     // load 4 elements from a
        float32x4_t vb = vld1q_f32(b + i);     // load 4 elements from b
        sum_vec = vfmaq_f32(sum_vec, va, vb);   // sum_vec += va * vb
    }

    // Reduce the 4 accumulators to a single scalar
    float sum = vaddvq_f32(sum_vec);

    // Handle remaining elements (if n is not a multiple of 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 代码中的通用模式：处理 vectorised block中的大部分内容，然后处理标量代码中的剩余部分。
为什么 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);
    }

    // Scalar cleanup
    for (; i < n; i++) {
        output[i] = input[i] > 0 ? input[i] : 0;
    }
}

vmaxq_f32 计算两个向量的元素最大值。由于 vector 全部为零，因此这正是 ReLU。没有branch，没有比较——只有一条指令。

I8MM：整数矩阵乘法¶

I8MM（Int8 矩阵乘法）是 ARMv8.6 extension，它添加了用于 INT8 矩阵乘法和 INT32 累加的专用指令，这正是量化 ML inference 所需要的。
关键指令是SMMLA（有符号矩阵乘法累加）：它采用两个 8×2 blocks INT8 值，并将结果累加到 2×2 block INT32 中：

#include <arm_neon.h>

// I8MM: multiply two 8-element INT8 vectors, accumulate into 4 INT32 results
// This computes a 2x2 tile of the output matrix from 2x8 x 8x2 input tiles
void matmul_i8mm_tile(const int8_t* A, const int8_t* B, int32_t* C) {
    // Load 8 bytes from A (2 rows of 4 elements, packed)
    int8x16_t va = vld1q_s8(A);   // 16 bytes = 2 rows × 8 elements
    int8x16_t vb = vld1q_s8(B);   // 16 bytes = 2 rows × 8 elements

    // Load existing accumulator (2x2 = 4 int32 values)
    int32x4_t acc = vld1q_s32(C);

    // I8MM instruction: acc += A_tile × B_tile^T
    // Computes 2×2 output from 2×8 × 8×2 inputs
    acc = vmmlaq_s32(acc, va, vb);  // THE I8MM instruction

    vst1q_s32(C, acc);
}

为什么 I8MM 很重要：如果没有 I8MM，NEON 上的 INT8 matmul 需要加宽乘法 (vmull)，然后是成对加法 — 每个输出元素有多个指令。使用 I8MM，硬件可以在一条指令中执行 8 元素点积 (2×8 × 8×2 = 2×2)。对于 INT8 inference 工作负载，这比普通 NEON 快 4-8 倍。
可用性：Apple M1+（所有 Apple Silicon）、ARM Cortex-A510/A710/X2+ (ARMv9)、AWS Graviton3+。检查 #ifdef __ARM_FEATURE_MATMUL_INT8。
对于 ML inference：在 ARM 服务器（Graviton）或 Apple Silicon 上运行的 INT8 量化 models（第 18 章）从 I8MM 中受益匪浅。 ONNX Runtime 和 llama.cpp 等框架在运行时检测 I8MM 并自动使用优化的 kernels。

SME 和 SME2：可扩展矩阵扩展¶

SME（可扩展矩阵扩展）是 ARM 对 Intel AMX 和 NVIDIA 张量核心的回答：用于矩阵运算的专用硬件。 SME2 (ARMv9.2) 进一步扩展了它。
SME 推出 ZA 瓦片 registers：存储在硬件中的 2D 矩阵，最多 SVL×SVL 字节（其中 SVL 是流 vector 长度，通常每维 128-512 位）。与 NEON（1D 矢量）甚至 SVE（1D 可缩放矢量）不同，SME 本身在 2D 切片 上运行。
model的编程有两种模式：
- 普通模式：标准ARM执行（NEON、SVE照常工作）。
- 流式SVE模式：通过smstart输入，启用SME指令。 SVE 指令也可以在此模式下工作，但可能使用不同的 register 宽度。

#include <arm_sme.h>

// SME2: outer product accumulation for matrix multiply
// Accumulates A_col × B_row into the ZA tile register
void sme2_matmul_outer(const float* A_col, const float* B_row, int K) {
    // Enter streaming mode
    // smstart;  // (done via compiler intrinsic or inline asm)

    // Zero the ZA tile accumulator
    svzero_za();

    for (int k = 0; k < K; k++) {
        // Load a column of A and a row of B into SVE registers
        svfloat32_t a = svld1_f32(svptrue_b32(), &A_col[k * SVL]);
        svfloat32_t b = svld1_f32(svptrue_b32(), &B_row[k * SVL]);

        // Outer product: ZA += a × b^T
        // This accumulates an SVL×SVL tile in one instruction
        svmopa_za32_f32_m(0, svptrue_b32(), svptrue_b32(), a, b);
    }

    // Store the ZA tile to memory
    // svst1_za(...);

    // Exit streaming mode
    // smstop;
}

关键概念：
- svmopa（外积累加）：核心SME指令。它计算两个向量的完整外积并累积到 ZA 瓦片中。对于 SVL=512 位（16 个浮点），这是一个 16×16 外积 — 一条指令中的 256 个 FMA 操作。
- ZA 磁贴：在流模式下跨指令保持不变。您将多个外积（每 K 次迭代一个）累积到同一个图block中，从而构建一个完整的矩阵乘法图block。
- 流模式：SME指令仅在流模式下工作。进入/退出流模式的开销意味着 SME 最适合持续矩阵计算，而不是短突发。
SME2 添加：多 vector 操作（同时处理 2 或 4 个 SVE 向量）、附加图block操作以及改进与正常模式的集成。
可用性：ARM Neoverse V2 (AWS Graviton4)，一些即将推出的移动芯片。尚未在 Apple Silicon 上（截至 2026 年）。 SME 仍处于早期阶段——大多数 ML 框架还没有针对 SME 优化的 kernels。
进展：NEON（128 位向量，按元素）→ I8MM（INT8 矩阵block）→ SVE（可扩展向量）→ SME（可扩展 2D 矩阵block）。每一代都更接近硬件中的本机矩阵运算。

SVE 和 SVE2：可扩展向量扩展¶

NEON 具有固定的 128 位宽度。 SVE（可扩展向量扩展）引入了vector 长度不可知（VLA）编程：您编写一次代码，它就可以在任何 vector 宽度（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);  // predicate: which lanes are active

    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();  // advance by the hardware vector width (in 32-bit elements)
        pred = svwhilelt_b32(i, n);
    }
}

谓词 registers (svbool_t) 替换标量清理循环。每个通道都有一个谓词位：活动通道参与，非活动通道被屏蔽。 svwhilelt_b32(i, n) 指令创建一个谓词，其中对应于 i, i+1, ..., n-1 的通道处于活动状态。这会自动处理尾部。
svcntw() 在运行时返回每个 vector register 的 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/等）用于后台任务。调度程序将 threads 分配给适当的核心类型。
AMX（Apple Matrix eXtensions）：专用矩阵乘法单元，与 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 写入的相同内存。这消除了机器学习工作负载的主要瓶颈。
神经引擎：16 核专用机器学习加速器。对 INT8 inference 执行约 30 TOPS（每秒万亿次操作）。由 Core ML 用于设备上的 inference。
对于Apple Silicon上的ML：使用MLX（Apple的ML框架），它是为统一内存架构而设计的。 PyTorch 还具有 MPS（金属性能着色器）后端支持，尽管它不如 CUDA 成熟。

自动矢量化¶

写 SIMD intrinsics 很乏味。 编译器可以自动矢量化您的代码吗？
是的，但有注意事项。现代编译器（GCC、Clang）可以自动向量化简单循环：

// The compiler CAN auto-vectorise this (with -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 tells the compiler: a, b, c point to non-overlapping memory
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）。
检查自动矢量化：使用编译器标志查看 vectorised 是什么：

# GCC: show vectorisation decisions
g++ -O3 -march=native -fopt-info-vec-optimized code.cpp

# Clang: show vectorisation report
clang++ -O3 -march=native -Rpass=loop-vectorize code.cpp

何时使用 intrinsics 与自动向量化：从干净的 C++ 和编译器优化开始。如果编译器对你的循环进行向量化，那就太好了。如果性能仍然不够，请检查编译器的向量化报告以了解原因，然后才为关键的内部循环编写 intrinsics。过早的 intrinsics 会使代码不可读且无法保证利益。

编码任务（在 ARM — Mac M 系列或 Linux aarch64 上使用 g++ 或 clang++ 进行编译）¶

编写标量点积和 NEON-vectorised 点积。对两者进行基准测试并测量加速比。

// task1_neon_dot.cpp
// Compile (Mac/ARM Linux): clang++ -O3 -o task1 task1_neon_dot.cpp
// Note: NEON is enabled by default on AArch64, no special flags needed

#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);

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

    // Benchmark scalar
    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;

    // Benchmark 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-max-查找。通过不同的操作练习加载→计算→存储模式。

// task2_neon_ops.cpp
// Compile: 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;
}

将自动 vectorised 代码与手写的 NEON intrinsics 进行比较。使用 -fopt-info-vec (GCC) 或 -Rpass=loop-vectorize (Clang) 进行编译，看看编译器做了什么。

// task3_auto_vs_manual.cpp
// Compile: clang++ -O3 -Rpass=loop-vectorize -o task3 task3_auto_vs_manual.cpp
//    (or): g++ -O3 -fopt-info-vec-optimized -o task3 task3_auto_vs_manual.cpp

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

// Let the compiler auto-vectorise
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];
    }
}

// Hand-written 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);  // warm up
        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");
    // They should be very close — the compiler auto-vectorises this simple loop well
    return 0;
}