第16章 §02 · ARM 与 NEON · 讲解稿¶

篇幅: ~4000 字 / 阅读 10 分钟

一 · 一句话总览¶

ARM = 手机/嵌入式之王, NEON 是其 128-bit SIMD 扩展, Apple Silicon (M1/M2/M3/M4/M5) 全部基于 ARM, 让 Mac 跑 ML 推理能与 NVIDIA GPU 掰手腕。

二 · ARM 架构概览¶

2.1 ARM vs x86 哲学¶

维度	ARM	x86
指令集	RISC (精简)	CISC (复杂)
功耗	极低 (1-15W)	中高 (45-250W)
性能/瓦	极佳	中
生态	手机/嵌入式/Mac	服务器/PC
SIMD	NEON (128) / SVE (可变)	SSE (128) / AVX (256/512)

2.2 Apple Silicon 崛起¶

Apple M-series (2020-2026): - M1 (2020): 8 核 CPU + 8 核 GPU, 16 核 NPU - M2 (2022): 24 GB 统一内存 - M3 (2023): 3nm 工艺 - M4 (2024): 38 TOPS NPU - M5 (2026): 60 TOPS NPU, 256-bit 内存

统一内存架构 (UMA): CPU + GPU + NPU 共享同一内存, 无需数据拷贝, LLM 推理极快。

2.3 移动 ARM¶

Snapdragon 8 Gen 3: Hexagon NPU 45 TOPS
联发科天玑 9300: APU 20 TOPS
Google Tensor G4: TPU 边缘版

三 · NEON 指令集¶

3.1 NEON 寄存器¶

32 个 128-bit 寄存器 (D0-D31 也可作 64-bit)
同一寄存器可看作:
4 × 32-bit float
8 × 16-bit float
16 × 8-bit int

3.2 NEON intrinsics (C)¶

#include <arm_neon.h>

// 4 路 float 加法
float32x4_t a = vld1q_f32(p1);   // 加载 4 个 float
float32x4_t b = vld1q_f32(p2);
float32x4_t c = vaddq_f32(a, b); // 4 个加法并行
vst1q_f32(out, c);                // 存储

3.3 NEON 实战: Softmax¶

void softmax_neon(const float *in, float *out, int n) {
    // 1. 找 max
    float32x4_t vmax = vld1q_f32(in);
    for (int i = 4; i < n; i += 4) {
        vmax = vmaxq_f32(vmax, vld1q_f32(in + i));
    }
    // 水平 max (4 个数归约)
    float max_val = vmaxvq_f32(vmax);

    // 2. exp & sum
    float32x4_t vsum = vdupq_n_f32(0);
    for (int i = 0; i < n; i += 4) {
        float32x4_t v = vsubq_f32(vld1q_f32(in + i), vdupq_n_f32(max_val));
        v = exp_ps(v);  // 需要实现向量 exp
        vsum = vaddq_f32(vsum, v);
        vst1q_f32(out + i, v);
    }
    float sum = vaddvq_f32(vsum);

    // 3. 除
    float inv_sum = 1.0f / sum;
    for (int i = 0; i < n; i += 4) {
        vst1q_f32(out + i, vmulq_f32(vld1q_f32(out + i), vdupq_n_f32(inv_sum)));
    }
}

3.4 NEON 在 ML 推理中的角色¶

Core ML (Apple): 自动使用 NEON + ANE (Apple Neural Engine)
XNNPACK (Google): 移动端推理引擎
ONNX Runtime Mobile: 自动 NEON 优化

四 · ARM Scalable Matrix Extension (SME)¶

4.1 SME 是什么¶

SME (Armv9): 矩阵扩展, 类似 NVIDIA Tensor Core - 矩阵 tile 操作 - 加速 Transformer / CNN - Apple M4+ 已支持 (改名为 AMX)

4.2 SME 编程¶

// SME 伪代码
void matmul_sme(float *A, float *B, float *C, int M, int N, int K) {
    // 启动 SME 模式
    sme_tile_zero(C_tile, sizeof(C_tile) / sizeof(float32_t));

    for (int k = 0; k < K; k += SME_TILE_K) {
        sme_load_a(A_tile, A + ...);
        sme_load_b(B_tile, B + ...);
        sme_mla(C_tile, A_tile, B_tile);  // 矩阵乘累加
    }

    sme_store(C, C_tile, ...);
}

4.3 SME 性能¶

理论: 比 NEON 快 4-8×
实测: Llama 2 7B INT4 推理, 比纯 NEON 快 3-5×

五 · Apple Neural Engine (ANE)¶

5.1 ANE 概览¶

专用 NPU (Neural Processing Unit)
矩阵乘 + 卷积优化
能效: 比 GPU 高 10×

5.2 性能演进¶

芯片	NPU TOPS	内存
A14 (2020)	11	6 GB
M1 (2020)	11	16 GB
M2 (2022)	15.8	24 GB
M3 (2023)	18	36 GB
M4 (2024)	38	64 GB
M5 (2026)	60	256-bit

5.3 ANE 应用¶

Core ML 自动编译到 ANE
实时语音转文字 (Whisper)
图像风格迁移 (Stable Diffusion)
LLM 推理 (Apple Intelligence / Private Cloud Compute)

六 · 一句话总结¶

ARM + NEON + ANE = 移动 + 端侧 ML 的三件套。M5 60 TOPS NPU + 256-bit UMA, 让 Mac/手机跑 7B 模型实时。未来 5 年 LLM 端侧推理主战场。