跳转至

# 为什么是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;
       }