为什么是C++以及ML框架如何工作

本书中每一次 jnp.matmul、每一次 torch.nn.Linear、每一次 np.dot 调用，底层都在执行C++和CUDA代码。本文档揭开帷幕：为何ML框架采用这种架构，面向Python工程师的C++快速入门，何时编写自定义C++核函数，以及如何将其绑定到Python——这是连接你所写代码与所运行硬件之间的桥梁。

• 你花了15章写Python。你导入了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丰富（笔记本、可视化、数据）	系统丰富（操作系统、驱动、引擎）

• 核心见解：每种语言发挥其优势。Python处理人力生产力重要的事务（实验设计、超参数调优、数据探索）。C++处理机器性能重要的事务（矩阵乘法、卷积、注意力核函数）。

• 一次矩阵乘法 jnp.matmul(A, B)，其中 \(A\) 为 \(4096 \times 4096\)，执行约1370亿次浮点运算。在纯Python（嵌套循环）中需要约30分钟。在使用AVX-512 SIMD和多线程优化后的C++中，只需约10毫秒。差距达180,000倍。再多的Python技巧也无法弥合这一鸿沟。

ML框架的结构

• 每个主流ML框架都遵循相同的架构：

用户代码（Python）
    ↓
Python API层（torch.nn、jax.numpy、numpy）
    ↓
调度/JIT编译器（torch.compile、XLA、NumPy调度）
    ↓
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函数，该函数调用BLAS（基本线性代数子程序），通常是Intel MKL或OpenBLAS。BLAS是手工优化的C和Fortran代码，使用SIMD指令、缓存感知的内存访问模式和多线程。数十年优化致力于让矩阵乘法更快。

• NumPy仅支持CPU，不使用GPU。但在CPU上，它极其快速，因为它委托给可用的最佳BLAS实现。

PyTorch

• PyTorch的计算引擎是ATen（张量库），用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张量对象中并返回。

• torch.compile（PyTorch 2.0+）更进一步：它追踪你的Python代码，构建计算图，并使用Triton（GPU）或C++/OpenMP（CPU）编译。编译后的代码融合操作，消除Python开销，可以比即时模式快2-5倍。

JAX

• JAX将Python函数编译为XLA（加速线性代数），Google的ML编译器。当你 jax.jit 一个函数时：

  1. JAX追踪函数，将操作捕获为XLA计算图（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* 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;
}

内存：栈与堆

// 栈分配：快速，自动生命周期（函数返回时释放）
float buffer[256];   // 栈上的256个浮点数

// 堆分配：手动，在函数外仍然存活
float* data = new float[n];   // 在堆上分配n个浮点数
// ... 使用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（第16章文件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, "自定义ReLU操作");
}

# 编译
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, "自定义GELU激活函数");
}

# 动态加载和编译
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自定义调用

• JAX使用XLA自定义调用。过程更为复杂（你需要向XLA注册一个C函数），但概念相同：编写C/C++，绑定，从Python调用。

• 对于大多数JAX用户，Pallas（在文件05中介绍）是更好的选择：它让你用类似Python的语法编写GPU核函数，由XLA编译，无需离开JAX生态系统。

大局观

• 本文解释了Python和硬件之间的层次。本章剩余文件将深入探讨：

  • 文件01：硬件本身（CPU架构、GPU架构、内存系统）
  • 文件02-03：CPU上的SIMD编程（ARM NEON、x86 AVX）——编写使用CPU向量单元的C++代码
  • 文件04：使用CUDA的GPU编程——编写在数千个GPU核心上运行的C++代码
  • 文件05：Triton、Pallas和更高级的GPU编程——编写编译为GPU核函数的Python代码

• 这种递进反映了抽象阶梯：C++内联函数（最低层、最多控制）→ CUDA（GPU专用）→ Triton/Pallas（Python风格、编译型）→ JAX/PyTorch（最高层、自动）。每一层以控制权换取便利性。理解较低层使你成为较高层的更好使用者。

编程任务（用g++或clang++编译）

1. 编写你的第一个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 << std::endl;
       std::cout << "时间: " << elapsed << " ms" << std::endl;
       std::cout << "元素数: " << N << std::endl;
       std::cout << "吞吐量: " << (N * sizeof(float)) / elapsed / 1e6 << " GB/s" << std::endl;
   
       return 0;
   }
   

2. 编写一个C++函数在数组上计算ReLU，然后使用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"匹配: {np.allclose(y_cpp, y_np)}")
   

3. 编写一个C++程序，演示为何内存布局很重要。比较行优先与列优先访问模式并测量性能差异。

   // 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_ms << " ms" << std::endl;
       std::cout << "列优先（缓存不友好）: " << col_ms << " ms" << std::endl;
       std::cout << "减速比: " << col_ms / row_ms << "x" << std::endl;
       std::cout << "（两个和: " << sum_row << ", " << sum_col << "）" << std::endl;
   
       return 0;
   }