第 16 章 SIMD 与 GPU 编程

第 16 章 SIMD 与 GPU 编程

理解硬件，才能榨干性能。本章带你从「调包」走向「写高性能内核」——拆开 PyTorch、JAX、NumPy 的黑盒，看清每一行 Python 背后的 C++、SIMD 与 GPU 代码，并亲手为 ARM、x86、GPU、TPU 写出真正跑得快的加速代码。

本章简介

到目前为止，你大部分时间都在写 Python：调 JAX、写 PyTorch、跑 NumPy。但你写下的每一次 jnp.matmul、每一个 torch.nn.Linear、每一次 np.dot，背后真正干活的都是 C++ 与 CUDA 代码。Python 只负责构造「计算描述」，真正跑在硬件上的，是 BLAS、cuBLAS、XLA 这些经过几十年打磨的高性能内核。

本章要回答一个核心问题：这些内核为什么快？你又该如何写出自己的内核？ 我们会从最底层的硬件讲起——摩尔定律的终结、SIMD 的本质、屋顶线模型——然后逐层往上：用 ARM NEON 与 x86 AVX 的 C++ 内联函数写 CPU 向量化代码；用 CUDA C++ 写运行在数千个核心上的 GPU 内核；用 Triton、Pallas 用 Python 风格写 GPU/TPU 内核；再到 RISC-V 与嵌入式设备上的 TinyML；最后用 Vulkan 和 WebGPU 做跨平台 GPU 计算。

这是一本直觉优先教程里最「硬核」的一章。但理解了它，你才能在 DeepMind、OpenAI 这类岗位上，真正读懂一个 Flash Attention、一个 fused kernel、一个量化推理引擎为什么这么写——并在性能不够时，知道该从哪里下手。

本章小节

• 00. 为什么用 C++，以及 ML 框架是如何工作的：揭开 Python 前端 + C++ 后端的双语言架构，讲清 NumPy、PyTorch（ATen）、JAX（XLA）的内部结构，并为 Python 工程师补齐 C++ 基础与 pybind11 绑定。
• 01. 硬件基础：摩尔定律为什么终结、并行为什么取代了主频、现代 CPU 的超标量与乱序执行、SIMD 的本质、用来推理性能的屋顶线模型，以及 x86/ARM/Apple Silicon/RISC-V/GPU/TPU 的全景。
• 02. ARM 与 NEON：ARM 架构基础、NEON 128 位 SIMD 内联函数、I8MM 整数矩阵乘、SME/SME2 可扩展矩阵扩展、SVE/SVE2、Apple Silicon 的 AMX 与统一内存、自动向量化，以及生产级范例 Cactus 引擎。
• 03. x86 与 AVX：x86 SIMD 演进史、AVX2 内联函数编程、AVX-512 与掩码寄存器、Intel AMX 矩阵乘硬件、内存对齐、性能陷阱（寄存器压力、数据依赖、循环展开）与性能剖析。
• 04. GPU 架构与 CUDA：GPU 与 CPU 的设计哲学差异、GPU 内存层级、CUDA 编程模型（Grid/Block/Thread）、warp 与 SIMT、内存合并、共享内存分块、流与并发、高级优化（AoS/SoA、内核融合、混合精度）与 NVIDIA 各代 GPU。
• 05. Triton、TPU 与 Pallas：用 Python 写 GPU 内核的 Triton、Flash Attention 案例、TPU 的脉动阵列架构、JAX/Pallas 编程 TPU，以及如何为不同工作负载挑对工具。
• 06. RISC-V 与嵌入式系统：开源 RISC-V 的哲学、V 向量扩展、微控制器上的 TinyML、TFLM、面向边缘的模型优化（量化、剪枝、蒸馏）、AI 加速器中的 RISC-V 控制核心，以及边缘部署的功耗/延迟/隐私约束。
• 07. Vulkan Compute 与跨平台 GPU：能在所有平台跑的 Vulkan、计算流水线、GLSL 计算着色器、完整 C++ 搭建、共享内存与同步、Kompute、浏览器里的 WebGPU，以及何时该用 Vulkan。

学习路径

• 前置知识：建议先读第 13 章「计算机体系结构」（理解 CPU 流水线、缓存层级、浮点数表示、bfloat16）和第 15 章「工程实践」（理解性能剖析、基准测试的基本方法论）。第 6 章里关于数据并行与混合精度训练的内容也会对本章有帮助。
• 后续章节：第 17 章「AI 推理」会大量用到本章的技术——量化内核、Flash Attention、Tensor Core、KV-cache 优化都建立在你在这里学到的 SIMD 与 GPU 编程之上。本章是你读懂任何高性能推理引擎（vLLM、TensorRT、llama.cpp、ONNX Runtime）的钥匙。

关键概念速览

• SIMD（Single Instruction, Multiple Data，单指令多数据）：一条指令同时处理多个数据元素，是 CPU 并行加速的核心机制，体现在 ARM NEON、x86 AVX/AVX-512 等指令集中。
• 屋顶线模型（roofline model）：用「峰值算力」和「峰值内存带宽」两个硬件极限，判断一段代码是算力受限还是访存受限，从而决定优化方向。
• warp（扭曲）：GPU 上 32 个线程为一组、同步执行同一条指令的基本调度单位；warp 分歧会让性能减半。
• 共享内存分块（tiling）：把一块数据从慢的全局内存加载到快的片上共享内存，反复复用，是写出快速 GPU 内核的最重要模式。
• 内核融合（kernel fusion）：把多个操作合成一个内核，避免把中间结果写回内存再读出来，对访存受限的操作尤其关键。
• Tensor Core / MXU：GPU 和 TPU 上专用的矩阵乘单元，一条指令完成一次小矩阵乘，是混合精度训练和推理高性能的硬件根基。