为什么用 C++ 与 ML 框架的工作原理¶

一句话总结：你写的每一行 Python ML 代码，底层跑的其实都是 C++ 与 CUDA；本节掀开这块幕布，讲清楚框架为何这么设计，以及何时该自己下场写内核。

🗺️ 本章导览¶

本章覆盖单指令多数据 (Single Instruction Multiple Data, SIMD) 与图形处理器 (Graphics Processing Unit, GPU) 编程：从 C++ 与框架 → 硬件 → ARM → x86 → GPU/统一计算设备架构 (Compute Unified Device Architecture, CUDA) → Triton/张量处理单元 (Tensor Processing Unit, TPU) → RISC-V → Vulkan
00. 为什么用 C++ 与 ML 框架的工作原理 (本节)：为什么机器学习 (Machine Learning, ML) 底层用 C++，框架又是如何组织起来的
01. 硬件基础：中央处理器 (Central Processing Unit, CPU) / GPU / 内存 / 存储层级
02. ARM 与 NEON：高级精简指令集机器 (Advanced RISC Machine, ARM) SIMD 指令集、I8MM、SME2
03. x86 与 AVX: x86 与高级向量扩展 (Advanced Vector Extensions, AVX) 指令集
04. GPU 架构与 CUDA: GPU 硬件模型、CUDA 编程
05. Triton、TPU 与 Pallas：高层 GPU 编程、TPU 生态
06. RISC-V 与嵌入式系统：开源指令集架构、嵌入式 AI
07. Vulkan Compute 与跨平台 GPU: WebGPU 生态

本书里每一个 jnp.matmul、每一个 torch.nn.Linear、每一个 np.dot 调用，背后跑的都是 C++ 与 CUDA 代码。这一节把幕布拉开：为什么 ML 框架要这么造，面向 Python 工程师的 C++ 速成，什么时候该自己写 C++ 内核，以及怎么把它绑回 Python. 这就是你写的代码与它真正跑在的硬件之间的那座桥.

你已经写了 15 章的 Python. 你 import JAX，调 jax.grad，跑训练循环，搭模型。整个体验都像 Python. 但真相是：几乎所有真正的计算都不是在 Python 里发生的.
当你在 PyTorch 里写 output = model(input)，或者在 JAX 里写 output = jnp.matmul(W, x), Python 其实几乎什么都没做。它只是构造一份"要算什么"的描述 (一张算子图)，然后把这张图交给 C++/CUDA 后端去干真正的活。 Python 是方向盘，C++ 才是发动机。

为什么前端是 Python，后端是 C++¶

这种双语言架构之所以存在，是因为 Python 和 C++ 恰好擅长相反的事:

	Python	C++
开发速度	快 (动态类型、有 REPL、不用编译)	慢 (静态类型、头文件、编译时间长)
执行速度	比 C 慢 100 倍左右 (解释执行，有 GIL)	接近硬件极限 (编译后无额外开销)
内存控制	自动 (垃圾回收 GC)，布局无法控制	手动，对每一个字节都能精确控制
硬件访问	没有 (访问不到 SIMD、GPU、自定义内存)	完全开放 (内建函数、CUDA、内联汇编)
生态	ML 生态丰富 (notebook、可视化、数据)	系统方向丰富 (操作系统、驱动、引擎)

一句话：让每种语言干它擅长的事. Python 负责人效重要的那部分 (实验设计、超参数调优、数据探索), C++ 负责机器效率重要的那部分 (矩阵乘法、卷积、注意力内核).
一次矩阵乘 jnp.matmul(A, B)，其中 \(A\) 是 \(4096 \times 4096\)，大约要做 1370 亿次浮点运算。用纯 Python 的嵌套循环跑，大概要 30 分钟。用带 AVX-512 SIMD 和多线程的优化版 C++，大概只要 10 毫秒。这是 18 万倍的差距。再怎么在 Python 里耍花招都补不上这个坑。

ML 框架是怎么组织的¶

所有主流 ML 框架其实都遵循同一套架构:

用户代码 (Python)
    ↓
Python API 层 (torch.nn, jax.numpy, numpy)
    ↓
派发 / 即时编译器 (torch.compile, XLA, NumPy dispatch)
    ↓
C++ 内核库 (ATen/PyTorch, XLA, BLAS/LAPACK)
    ↓
硬件专属后端 (CUDA, cuDNN, MKL, oneDNN, Metal)
    ↓
硬件 (CPU SIMD 单元, GPU 核心, TPU MXU)

NumPy¶

NumPy 的核心是用 C 写的。当你调 np.dot(A, B), Python 会调一个 C 函数，这个函数再去调基础线性代数子程序 (Basic Linear Algebra Subprograms, BLAS)，通常是 Intel 的 MKL 或者 OpenBLAS. BLAS 是一坨手工优化的 C 和 Fortran 代码，里面用了 SIMD 指令、缓存友好的内存访问模式和多线程。数十年的优化经验都堆在"怎么把矩阵乘算快"这件事上。
NumPy 只跑 CPU，不用 GPU. 但在 CPU 上，它极快，因为它会把活委托给当前平台上最好的 BLAS 实现。

PyTorch¶

PyTorch 的计算引擎叫 ATen (A Tensor Library，张量库)，用 C++ 写的。 ATen 实现了大约 2000 个张量算子 (add、matmul、conv2d、softmax……)，每个都带 CPU 和 CUDA 两套后端。
当你调 torch.matmul(A, B):
1. Python 把调用派发到对应的 ATen C++ 函数。
2. ATen 检查设备 (CPU 还是 CUDA) 和数据类型。
3. CPU 上：调 MKL 或 OpenBLAS. GPU 上：调 cuBLAS (NVIDIA 的 GPU 优化版 BLAS).
4. 结果被包一层 Python tensor 对象再返回。
torch.compile (PyTorch 2.0+) 更进一步：它会追踪你的 Python 代码，构建一张计算图，然后用 Triton (面向 GPU) 或 C++/OpenMP (面向 CPU) 编译这张图。编译后的代码会把算子融合起来、消掉 Python 开销，可以比 eager 模式快 2 到 5 倍。

JAX¶

JAX 把 Python 函数编译成 XLA (Accelerated Linear Algebra，加速线性代数) 代码，这是 Google 面向 ML 负载的编译器。当你对函数调 jax.jit:
1. JAX 追踪函数，把里面的算子记录成一张 XLA 计算图 (High Level Operations, HLO，高层算子).
2. XLA 优化这张图：融合算子、消除冗余计算、优化内存布局。
3. XLA 编译到目标后端：CPU (走 LLVM)、GPU (走 CUDA/PTX)，或 TPU (走 TPU 专属指令).
4. 编译后的代码直接跑在硬件上，Python 完全不掺和。
这就是为什么 jax.jit 这么关键：不用它，每个算子都要走一遍 Python↔C++ 的往返；用了它，整个函数就变成了一个编译好的内核。

面向 Python 工程师的 C++ 速成¶

你不需要成为 C++ 专家。你只需要懂到能读懂内核代码、能写点简单扩展、能听懂性能讨论就行。下面是最必要的一点点。

类型与变量¶

// C++ 必须显式声明类型 (跟 Python 不一样)
int count = 0;           // 32 位整数
float loss = 0.5f;       // 32 位浮点
double lr = 3e-4;        // 64 位浮点
bool training = true;    // 布尔

// 数组 (固定大小, 栈上分配)
float weights[1024];     // 1024 个 float, 在内存里连续

// 指针: 一个保存内存地址的变量
float* ptr = weights;    // ptr 指向 weights 的第一个元素
float val = ptr[42];     // 通过指针算术访问第 42 个元素
// ptr[42] 等价于 *(ptr + 42)

指针是与 Python 最大的观念差异。在 Python 里，一切都是引用，你从不需要想内存地址。在 C++ 里，指针给你对内存的直接访问 — 强大但危险 (悬空指针、缓冲区溢出).

函数¶

// 函数声明: 返回类型 名字(参数类型 参数名)
float relu(float x) {
    return x > 0.0f ? x : 0.0f;
}

// 按引用传参 (避免拷贝大对象)
void scale_vector(std::vector<float>& vec, float factor) {
    for (size_t i = 0; i < vec.size(); i++) {
        vec[i] *= factor;
    }
}

// const 引用: 只读, 不拷贝
float sum(const std::vector<float>& vec) {
    float total = 0.0f;
    for (float x : vec) {  // 基于范围的 for 循环 (类似 Python 的 for x in vec)
        total += x;
    }
    return total;
}

内存：栈 vs 堆¶

// 栈分配: 快, 生命周期自动 (函数返回时自动释放)
float buffer[256];   // 栈上分配 256 个 float

// 堆分配: 手动, 生命周期超出当前函数
float* data = new float[n];   // 堆上分配 n 个 float
// ... 使用 data ...
delete[] data;                 // 必须自己释放 (没有垃圾回收)

// 现代 C++: 智能指针 (自动清理, 类似 Python 的引用)
#include <memory>
auto data = std::make_unique<float[]>(n);  // 离开作用域时自动释放

关键法则：栈很快但很小 (通常 1 到 8 MB). 大数组 (张量、特征图) 必须放在堆上。 Python 里一切都在堆上，由 GC 帮你清理；C++ 里你要么自己管，要么用智能指针。

模板 (泛型)¶

// 一个能处理任意数值类型的函数
template <typename T>
T add(T a, T b) {
    return a + b;
}

add<float>(1.5f, 2.5f);   // 返回 4.0f
add<int>(3, 4);             // 返回 7

C++ 库 (比如 ATen) 就是靠模板写出一份能同时处理 float16、float32、float64 的代码，而不用把实现复制多份。

标准库必备¶

#include <vector>      // 动态数组 (类似 Python list)
#include <string>      // 字符串
#include <unordered_map>  // 哈希表 (类似 Python dict)
#include <algorithm>   // sort、find、transform 等
#include <cmath>       // 数学函数

std::vector<float> vec = {1.0f, 2.0f, 3.0f};
vec.push_back(4.0f);            // 追加
float first = vec[0];           // 索引
size_t len = vec.size();        // 长度

std::unordered_map<std::string, int> counts;
counts["hello"] = 5;            // 插入
if (counts.count("hello")) { }  // 判断存在

什么时候该自己写 C++ 内核¶

绝大多数 ML 工程师一辈子都不用写 C++. 框架内置的算子覆盖了 99% 的场景。只有在下面这些情况下才该考虑自己写 C++:
框架里根本没有你要的算子：一种新颖的激活函数、一种自定义的注意力模式、一种没法用现有算子组合出来的特殊损失函数。
为了性能融合算子：你的模型要做 relu(layernorm(matmul(x, W) + b)). 每个算子都要单独发起一次内核、读写一次内存、同步一次。一个融合内核可以一遍走完，避免内存反复往返。这可能带来 2 到 5 倍加速。
降低内存占用：一个自定义内核可以在不保存所有中间激活的前提下算出梯度 (相当于在内核层面做梯度检查点).
面向新硬件：一款新加速器 (比如 Cerebras、Groq) 可能还没有框架支持，你就得直接写内核。
对情况 1 和 2, Triton (本章第 05 节) 通常就够用，而且比直接写 CUDA C 容易得多。只有当 Triton 表达不了你要的东西时，才降到 CUDA C.

怎么把 C++ 绑回 Python¶

写 C++ 只是活干了一半，你还得能从 Python 里调它。

pybind11 (通用方案)¶

pybind11 能用极少的样板代码为 C++ 函数生成 Python 绑定:

// my_ops.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
namespace py = pybind11;

// 一个简单的自定义算子
py::array_t<float> custom_relu(py::array_t<float> input) {
    auto buf = input.request();
    float* ptr = static_cast<float*>(buf.ptr);
    size_t n = buf.size;

    auto result = py::array_t<float>(n);
    float* out = static_cast<float*>(result.request().ptr);

    for (size_t i = 0; i < n; i++) {
        out[i] = ptr[i] > 0 ? ptr[i] : 0;
    }
    return result;
}

PYBIND11_MODULE(my_ops, m) {
    m.def("custom_relu", &custom_relu, "Custom ReLU operation");
}

# 编译
pip install pybind11
c++ -O3 -shared -std=c++17 -fPIC $(python3 -m pybind11 --includes) my_ops.cpp -o my_ops$(python3-config --extension-suffix)

# 在 Python 里调用
import my_ops
import numpy as np

x = np.array([-1.0, 2.0, -3.0, 4.0], dtype=np.float32)
y = my_ops.custom_relu(x)
print(y)  # [0. 2. 0. 4.]

PyTorch C++ 扩展¶

PyTorch 提供了一条更顺的路来加自定义算子:

// custom_op.cpp
#include <torch/extension.h>

torch::Tensor custom_gelu(torch::Tensor x) {
    return x * 0.5 * (1.0 + torch::erf(x / std::sqrt(2.0)));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("custom_gelu", &custom_gelu, "Custom GELU activation");
}

# 现场加载与编译
from torch.utils.cpp_extension import load

custom_ops = load(
    name="custom_ops",
    sources=["custom_op.cpp"],
    extra_cflags=["-O3"],
)

x = torch.randn(1000)
y = custom_ops.custom_gelu(x)

torch.utils.cpp_extension.load 会一次性把 C++ 代码编译、生成动态库、加载成 Python 模块。这是在 PyTorch 里试玩自定义 C++ 算子最省事的路子。

JAX Custom Calls¶

JAX 用的是 XLA 的自定义调用机制。流程更绕一点 (你要把一个 C 函数注册到 XLA)，但概念一样：写 C/C++、绑定、从 Python 调。
对大多数 JAX 用户来说，Pallas (第 05 节) 是更好的选择：它让你用类似 Python 的语法写 GPU 内核，由 XLA 编译，完全不用离开 JAX 生态。

大局观¶

这一节讲的是 Python 与硬件之间那一层。本章后面的文件继续往下钻:
- 第 01 节：硬件本身 (CPU 架构、GPU 架构、内存系统)
- 第 02 到 03 节: CPU 上的 SIMD 编程 (ARM NEON、x86 AVX) — 你在这一层用 C++ 直接调 CPU 的向量单元
- 第 04 节：用 CUDA 做 GPU 编程 — 你在这一层写 C++ 让代码跑在成千上万个 GPU 核心上
- 第 05 节: Triton、Pallas 与更高层的 GPU 编程 — 你在这一层写 Python，由编译器帮你产出 GPU 内核
这条路径正对应抽象梯子：C++ 内建函数 (最底层，控制最强) → CUDA (GPU 专属) → Triton/Pallas (Python 风格，会被编译) → JAX/PyTorch (最高层，全自动). 每上一格，都是用控制换便利。懂了底下几层，你才是上层的更好用户。

编码任务 (用 g++ 或 clang++ 编译)¶

写你的第一个 C++ 程序。分配一个数组、填充它、算总和、量时间。这一步会带你走一遍编译、数组、指针和计时。

// task1_basics.cpp
// 编译: g++ -O3 -o task1 task1_basics.cpp
// 运行: ./task1

#include <iostream>
#include <chrono>
#include <vector>

int main() {
    const int N = 10'000'000;  // C++ 允许用 ' 作数字分隔符
    std::vector<float> data(N);

    // 填充数组
    for (int i = 0; i < N; i++) {
        data[i] = static_cast<float>(i) * 0.001f;
    }

    // 计算总和
    auto start = std::chrono::high_resolution_clock::now();
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        sum += data[i];
    }
    auto end = std::chrono::high_resolution_clock::now();
    double elapsed = std::chrono::duration<double, std::milli>(end - start).count();

    std::cout << "Sum: " << sum << std::endl;
    std::cout << "Time: " << elapsed << " ms" << std::endl;
    std::cout << "Elements: " << N << std::endl;
    std::cout << "Throughput: " << (N * sizeof(float)) / elapsed / 1e6 << " GB/s" << std::endl;

    return 0;
}

写一个在数组上算 ReLU 的 C++ 函数，然后用 pybind11 做一个 Python 绑定。从 Python 里调它，跟 NumPy 比一下速度。

// task2_relu.cpp
// 编译: c++ -O3 -shared -std=c++17 -fPIC $(python3 -m pybind11 --includes) \
//          task2_relu.cpp -o my_relu$(python3-config --extension-suffix)

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
namespace py = pybind11;

py::array_t<float> cpp_relu(py::array_t<float> input) {
    auto buf = input.request();
    float* ptr = static_cast<float*>(buf.ptr);
    int n = buf.size;

    auto result = py::array_t<float>(n);
    float* out = static_cast<float*>(result.request().ptr);

    for (int i = 0; i < n; i++) {
        out[i] = ptr[i] > 0.0f ? ptr[i] : 0.0f;
    }
    return result;
}

PYBIND11_MODULE(my_relu, m) {
    m.def("relu", &cpp_relu, "C++ ReLU");
}

# test_relu.py — 先编译好上面的 C++ 模块再运行
import numpy as np
import time
import my_relu  # 编译好的 C++ 模块

x = np.random.randn(10_000_000).astype(np.float32)

# C++ ReLU
start = time.time()
for _ in range(100):
    y_cpp = my_relu.relu(x)
cpp_time = (time.time() - start) / 100

# NumPy ReLU
start = time.time()
for _ in range(100):
    y_np = np.maximum(x, 0)
np_time = (time.time() - start) / 100

print(f"C++ ReLU:   {cpp_time*1000:.2f} ms")
print(f"NumPy ReLU: {np_time*1000:.2f} ms")
print(f"Match: {np.allclose(y_cpp, y_np)}")

写一个 C++ 程序，演示内存布局为什么重要。对比行主序 (row-major) 与列主序 (column-major) 访问模式，量出性能差异。

// task3_layout.cpp
// 编译: g++ -O3 -o task3 task3_layout.cpp

#include <iostream>
#include <chrono>
#include <vector>

int main() {
    const int N = 4096;
    std::vector<float> matrix(N * N, 1.0f);

    // 行主序访问: 连续内存地址 (缓存友好)
    auto start = std::chrono::high_resolution_clock::now();
    float sum_row = 0.0f;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            sum_row += matrix[i * N + j];  // 步长为 1 的访问
        }
    }
    auto end = std::chrono::high_resolution_clock::now();
    double row_ms = std::chrono::duration<double, std::milli>(end - start).count();

    // 列主序访问: 步长为 N 的访问 (缓存不友好)
    start = std::chrono::high_resolution_clock::now();
    float sum_col = 0.0f;
    for (int j = 0; j < N; j++) {
        for (int i = 0; i < N; i++) {
            sum_col += matrix[i * N + j];  // 步长 N 的访问 (缓存缺失!)
        }
    }
    end = std::chrono::high_resolution_clock::now();
    double col_ms = std::chrono::duration<double, std::milli>(end - start).count();

    std::cout << "Row-major (cache-friendly): " << row_ms << " ms" << std::endl;
    std::cout << "Col-major (cache-hostile):  " << col_ms << " ms" << std::endl;
    std::cout << "Slowdown: " << col_ms / row_ms << "x" << std::endl;
    std::cout << "(Both sums: " << sum_row << ", " << sum_col << ")" << std::endl;

    return 0;
}

📌 本节要点¶

双语言架构: Python 前端负责开发效率，C++ 后端负责执行效率，各干各擅长的事。
Python 只是描述: torch.matmul 之类的调用只在 Python 里构建算子图，真正的计算全在 C++/CUDA 里。
NumPy 走 BLAS: NumPy 的核在 C 里，靠 MKL/OpenBLAS 把矩阵乘算到 CPU 极限。
PyTorch 靠 ATen：大约 2000 个算子，CPU 走 MKL、GPU 走 cuBLAS, torch.compile 再叠一层图融合。
JAX 靠 XLA: jax.jit 把整个函数编译成一个 XLA 内核，免掉每步的 Python↔C++ 往返。
C++ 核心概念：显式类型、指针、栈 vs 堆、模板、标准库，够读懂内核就行。
何时下场写 C++：缺算子、要融合、要压内存、面向新硬件；前两种情况优先考虑 Triton.
绑定回 Python: pybind11 是通用方案，PyTorch 有 cpp_extension.load 更省事，JAX 更推荐 Pallas.
抽象梯子：内建函数 → CUDA → Triton/Pallas → JAX/PyTorch，每上一格是"用控制换便利".
内存布局重要：行主序连续访问对缓存友好，列主序跨行访问会引发大量缓存缺失。