《深度学习框架PyTorch：入门与实战》第8章 CUDA 扩展与编译实战

本章将带你完整走一遍 PyTorch 的 C++ 扩展与 CUDA 扩展流程：从编写自定义前向/反向函数、用 pybind11 绑定 Python、setuptools 编译安装，到在 autograd.Function 与 nn.Module 中调用并验证正确性，同时剖析 CUDA 运行版本、编译版本与 cuDNN 之间的关系。读完本章，你将能自己动手为 PyTorch 写一个可编译、可安装、结果可验证的自定义算子扩展。

图 8-1 CUDA 中的 Grid、Block 与 Thread（来源：Chapter8/img/Grid.png）

一、为什么需要 C++ / CUDA 扩展

PyTorch 已经内置了非常丰富的算子，但某些场景下你仍然需要自定义操作，主要集中在两类情况：

• PyTorch 尚未支持该操作：需要自行实现新的算法或算子；
• PyTorch 中对该操作的支持不够高效：例如某个功能可以靠循环调用 PyTorch 函数实现，但效率很低，原因有二：
  • Python 调用本身开销大：每执行一次 Python 命令，都要经由 Python 解释器调用 CPython，再调用 C/C++ 扩展。当循环调用次数很多时，速度差异会非常明显；
  • 计算图优化受限：反向传播时 PyTorch 依靠 autograd 生成梯度计算规则，但该规则仍有优化空间。尤其当多次调用 Python 命令时，会生成非常复杂的计算图，严重增加系统开销。

从逻辑上讲，最直接的自定义方式是在 Python 层面改写 Function 和 Module——重写前向传播与反向传播（类似本书第 3 章中扩展 torch.autograd 的做法）。但这样做存在性能与逻辑上的缺陷：

• 运行时代码需要调用 Python 解释器，会拖慢程序速度；
• PyTorch 无法真正了解自定义类的算法整体逻辑，只能逐一优化其中每个独立的操作；
• PyTorch 必须按顺序依次执行这些操作，而每次调用都可能存在启动 CUDA 内核的开销；当独立操作很多时，性能会显著下降。

因此，当代码需要被反复调用且对性能要求较高，或者需要与 C/C++ 对接时，就应当使用本章介绍的 C++ 扩展与 CUDA 扩展。

仓库配套源码位于 Chapter8/CppExtension 与 Chapter8/CUDAExtension，下文讲解的 y = x² + b 示例与仓库中实际的 Sigmoid 示例结构一致，可对照学习。

二、C++ 扩展实战（以 y = x² + b 为例）

普通的 C++ 扩展可以分为三步：

1. 定义并实现 C++ 函数，使用 pybind11 生成函数的 CPython 接口；
2. 使用 setuptools 编译安装；
3. 自定义 autograd.Function 调用该函数的 CPython 接口，随后便可在 nn.Module 中使用。

2.1 工程目录结构

|__ CppExtension
    |—— src
        |—— MyLinear.h
        |__ MyLinear.cpp
    |__ setup.py

本仓库实际提供的对应工程为 Chapter8/CppExtension，其目录结构为 src/MySigmoid.h、src/MySigmoid.cpp、setup.py、test.py，实现的自定义函数为 Sigmoid。

2.2 编写 C++ 头文件

头文件 MyLinear.h 声明了前向与反向两个函数。这里使用了 torch/torch.h 这个一站式头文件，它包含：

• ATen——张量计算接口库；
• pybind11——用于绑定 C++ 与 Python（使用前需在 Python 环境中执行 pip install pybind11）；
• 管理 ATen 与 pybind11 的其他文件。

// ./src/MyLinear.h
#include<torch/torch.h>
#include<vector>
#include<iostream>

// 前向传播 过程中用到了参数x和b
std::vector<at::Tensor> MyLinearForward(const at::Tensor& x, const at::Tensor& b);
// 反向传播 过程中用到了回传参数和参数x
std::vector<at::Tensor> MyLinearBackward(const at::Tensor& x, const at::Tensor& grad_out);

2.3 实现 C++ 函数并用 pybind11 绑定

C++ 扩展 API 不能像 autograd 一样自动计算梯度，因此需要手动完成反向传播的计算过程，再用 pybind11 把 C++ 函数绑定到 Python 上。宏 TORCH_EXTENSION_NAME 会被定义为 setup.py 脚本中设置的扩展名名称。

// ./src/MyLinear.cpp
#include "MyLinear.h" // 导入头文件

// 前向传播函数，注意Tensor的shape要一致
std::vector<at::Tensor> MyLinearForward(const at::Tensor& x, const at::Tensor& b){
    at::Tensor y = at::zeros(x.sizes());
    y = x * x + b;
    return {x, y};
}

// 反向传播函数，grad_out是上一级梯度回传的值，根据链式法则需要乘上去
std::vector<at::Tensor> MyLinearBackward(const at::Tensor& x, const at::Tensor& grad_out){
    at::Tensor grad_x = 2 * x * grad_out * at::ones(grad_out.sizes());
    at::Tensor grad_b = grad_out * at::ones(grad_out.sizes());
    return {grad_x, grad_b};
}

// 绑定C++函数到python上，TORCH_EXTENSION_NAME即为setup.py中setup定义的name
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward",  &MyLinearForward,  "mylinear forward" );
  m.def("backward", &MyLinearBackward, "mylinear backward");
}

2.4 用 setuptools 编译安装

利用 torch.utils.cpp_extension 中的 CppExtension 与 BuildExtension 完成编译和构建。构建完成后扩展被命名为 mylinear，之后便可在 Python 中 import 该模块。

# ./setup.py
from setuptools import setup
import os
from torch.utils.cpp_extension import BuildExtension, CppExtension

# 头文件目录
include_dirs = "./src"
# 源代码目录
source = ["./src/MyLinear.cpp"]

setup(
    name='mylinear',  # 模块名称，宏TORCH_EXTENSION_NAME的值
    version="0.1",
    ext_modules=[CppExtension('mylinear', sources=source, include_dirs=[include_dirs]),],
    cmdclass={'build_ext': BuildExtension}
)

在 setup.py 所在目录运行：

python setup.py install

完成后模块会安装到 Python 的 site-packages 中。官方文档特别强调编译版本问题：用于构建 C++ 扩展的编译器必须与 ABI 兼容，且必须与构建 PyTorch 时采用的编译器一致，即在 Linux 上需使用 4.9 及更高版本的 GCC。

本仓库中的实际构建脚本见 Chapter8/CppExtension/setup.py，它将模块命名为 mysigmoid，其余配置与上述示例完全一致。

2.5 在 Python 中调用并验证正确性

编写测试代码，将扩展的前向、反向封装为 autograd.Function，再放入 nn.Module，最后与纯 PyTorch 计算对比前向输出与梯度：

import torch
from torch.autograd import Function
import mylinear # 导入我们的扩展

# 将扩展的前向传播和反向传播封装为一个Function对象
class MyLinear(Function):
    @staticmethod
    def forward(ctx, x, b):
        t_x, t_y = mylinear.forward(x, b)
        vars = [t_x]
        ctx.save_for_backward(*vars) # 必须用*解包后传入
        return t_y

    @staticmethod
    def backward(ctx, grad_out):
        grad_x, grad_b = mylinear.backward(*ctx.saved_tensors, grad_out)
        # 必须将mylinear.backward返回值解耦后分别返回，不然是一个对象
        return grad_x, grad_b

# 构建用于测试的示例类Test
class Test(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, b):
        return MyLinear.apply(x, b)

# 测试
def main():
    x1 = torch.tensor([1.,2.,3.], requires_grad=True)
    b1 = torch.tensor([1.,2.,3.], requires_grad=True)
    x2 = x1.clone().detach().requires_grad_()
    b2 = b1.clone().detach().requires_grad_()

    # C++ Extension
    model = Test()
    y1 = model(x1, b1)
    y1.sum().backward()

    # Python接口
    y2 = x2 ** 2 + b2
    y2.sum().backward()

    assert y1.equal(y2)
    assert x1.grad.equal(x2.grad)
    assert b1.grad.equal(b2.grad)

if __name__ == '__main__':
    main()

运行测试函数时如果没有抛出异常，说明自定义的 C++ 扩展与直接使用 PyTorch 反向传播的结果完全一致。

2.6 仓库补充：Sigmoid 扩展与性能对比实现

仓库中的真实实现是 Sigmoid 算子，见 Chapter8/CppExtension/src/MySigmoid.cpp。其前向为 f(x) = e⁻ˣ / (1 - e⁻ˣ)，反向为 f'(x) = -f(x)·(f(x)+1)·grad_out，同时给出了逐元素 for 循环与向量化写法两种实现（源码中用注释标出，可互相替换），便于读者体会 C++ 循环与 ATen 向量化运算的差异。

仓库配套的 Chapter8/CppExtension/test.py 除了做正确性断言（fx1.equal(fx2) 与 x1.grad.allclose(x2.grad)）外，还内置了性能对比函数 compare_pytorch_with_cpp()：

• 使用 torch.set_num_threads(1) 固定线程数，保证对比环境一致；
• 对 (1280, 1280) 的随机张量分别计时 C++ 扩展 forward/backward、PyTorch 双层 for 循环、PyTorch 向量化运算，打印各自耗时。

这正好印证了 8.1.3 节的结论：C++ 扩展省去了大量 Python 上层调用，速度明显快于 Python 原生循环。

三、CUDA 扩展实战

当扩展中包含大量逐点运算和矩阵运算时，C++ 扩展的性能依旧有限。此时可以自定义 CUDA 内核，将逐点运算和矩阵运算放进内核中进行融合与并行化，进一步压榨 GPU 性能。

编写 CUDA 扩展的流程与 C++ 扩展类似：

1. 编写一个 C++ 文件，定义 Python 中会调用的函数，用 pybind11 绑定，并将调用转发给 CUDA 函数；
2. 编写一个以 .cu 为后缀的 CUDA 文件，自定义实际的 CUDA 内核，内部使用 CUDA 语法；
3. 用 CUDAExtension（而非 CppExtension）编译安装；
4. 在 Python 中通过 autograd.Function 调用。

3.1 编写 C++ 绑定层（host 侧）

这部分与 C++ 扩展几乎一致（实现功能相同）：

// ./src/linear_cuda.cpp
#include <torch/torch.h>
#include <vector>

// 前向传播
std::vector<at::Tensor> linear_cuda_forward(const at::Tensor& x, const at::Tensor& b){
    at::Tensor y = at::zeros_like(x);
    y = x * x + b;
    return {x, y};
}

// 反向传播
std::vector<at::Tensor> linear_cuda_backward(const at::Tensor& x, const at::Tensor& grad_y){
    at::Tensor grad_x = 2 * x * grad_y * at::ones_like(grad_y);
    at::Tensor grad_b = grad_y * at::ones_like(grad_y);
    return {grad_x, grad_b};
}

// 绑定C++函数到python上
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward",  &linear_cuda_forward,  "linear cuda forward");
  m.def("backward", &linear_cuda_backward, "linear cuda backward");
}

3.2 CUDA 内核的基础概念

使用 CUDA 即意味着获得了 GPU 的控制权。由于 CUDA 中有非常多的核（线程）在执行运算，调用函数时需要指明使用哪些线程。CUDA 将一个 kernel 分为三级：Grid、Block 和 Thread。Grid 和 Block 都是 dim3 类型的变量，可看作无符号整数三元组 (x, y, z)；初始化时灵活地将部分维设为 1，即可轻松实现 1-dim、2-dim 与 3-dim 结构。

定义完 Grid 和 Block 后，每个线程的全局 ID 可用一组坐标唯一标识：(blockIdx, threadIdx)。其中 blockIdx 表示 Block 在 Grid 中的位置，threadIdx 表示线程在 Block 中的位置。例如图 8-1 中 Block(2,1) 里的 Thread(2,1) 可表示为：

threadIdx.x = 2
threadIdx.y = 1
blockIdx.x = 1
blockIdx.y = 2

通常还需把一个线程在 Block 中的位置从多维坐标转为一维形式（类似计算 reshape 后 Tensor 某元素对应的底层 Storage Offset）。对二维 Block (blockDim.x, blockDim.y)，线程 (x, y) 的一维 ID 为：

$$x + y \times blockDim.x$$

此外，CUDA 的多线程在逻辑层与物理层上存在区别：kernel 启动的所有线程构成一个 grid，再被分成多个 block；这些线程在逻辑上是并行的，物理上却不一定。GPU 硬件的核心组件是 SM（流式多处理器），包含 CUDA 核心、共享内存与寄存器。kernel 执行时，grid 中的线程块会被分配到 SM 上（一个 SM 可调度多个线程块，但一个线程块只能由一个 SM 调度）。由于 SM 同时处理的线程数有限，同一 kernel 下的所有线程物理上不一定同时并发。因此在编写核函数时，需要有效利用线程序号分配计算任务，尽量把计算量平均分给各个线程；若计算量超过线程数量，系统会循环地把任务分配给线程。正是这种线程组织结构，让 kernel 在矩阵运算上十分高效。

CUDA 中还有几种特殊的声明方式：

• __global__：异步模式，在 CPU 中调用、在 GPU 上执行，CPU 不会等待 kernel 执行完就继续下一步；
• __device__：直接从 GPU 调用、在 GPU 中执行，不可与 __global__ 同时使用；
• __host__：同步模式，从 CPU 上调用、在 CPU 中执行，一般可省略不写。

3.3 编写 CUDA 内核（device 侧）

// ./src/linear_cuda_kernel.cu （注意后缀名是.cu）
#include <ATen/ATen.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

// scalar_t是一个宏，特化的时候会传入具体的类型，下面调用的时候实例化为at::Tensor
// 定义了前向传播的cuda内核
template <typename scalar_t>
__global__ void linear_cuda_forward_kernel(
const scalar_t* x, const scalar_t* b, scalar_t* y) {
    const int index = blockIdx.x + blockDim.x * blockIdx.y;
    y[index] = x[index] * x[index] + b[index];
}

// 定义了反向传播的cuda内核
template <typename scalar_t>
__global__ void linear_cuda_backward_kernel(
const scalar_t* grad_y, const scalar_t* x, scalar_t* grad_x, scalar_t* grad_b ) {
    const int index = blockIdx.x + blockDim.x * blockIdx.y;
    grad_x[index] = 2 * x[index] * grad_y[index];
    grad_b[index] = grad_y[index];
}


at::Tensor linear_cuda_forward(at::Tensor x, at::Tensor b) {
    auto y = at::zeros_like(x);
    // forward kernel线程配置
    const dim3 blocks(x.size(0), x.size(1));
    const int threads = 4;
    // forward kernel调用，使用参数按照引用传递的匿名函数
    AT_DISPATCH_FLOATING_TYPES(x.type(), "linear_forward_cuda", ([&] {
        linear_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
        x.data<scalar_t>(),
        b.data<scalar_t>(),
        y.data<scalar_t>());
  }));
  return y;
}

at::Tensor linear_cuda_backward(at::Tensor grad_y, at::Tensor x) {
    auto grad_x = at::zeros_like(grad_y);
    auto grad_b = at::zeros_like(grad_y);
    // backward kernel线程配置
    const dim3 blocks(grad_y.size(0), grad_y.size(1));
    const int threads = 4;
    // backward kernel调用  使用参数按照引用传递的匿名函数
    AT_DISPATCH_FLOATING_TYPES(grad_y.type(), "linear_backward_cuda", ([&] {
    linear_cuda_backward_kernel<scalar_t><<<blocks, threads>>>(
        grad_y.data<scalar_t>(),
        x.data<scalar_t>(),
        grad_x.data<scalar_t>(),
        grad_b.data<scalar_t>());
  }));
  return grad_x, grad_b;
}

要点说明：

• AT_DISPATCH_FLOATING_TYPES 用于按 Tensor 的实际浮点类型（float/double/half 等）特化模板 kernel，其中 scalar_t 即特化后传入的具体类型；
• x.data<scalar_t>() 取得底层指针交给 kernel；
• <<<blocks, threads>>> 是 CUDA 的 kernel 启动配置语法，指定 Grid/Block 尺寸与线程数。

3.4 仓库补充：更规范的 Sigmoid CUDA 内核

仓库中的实际实现 Chapter8/CUDAExtension/src/MySigmoidKernel.cu 演示了更工程化的写法，值得对照学习：

• 一维扁平化索引 + 越界保护：

#define THREADS 1024

template <typename scalar_t>
__global__ void sigmoid_cuda_forward_kernel(scalar_t* x, scalar_t* fx, const int state_size) {
    const uint32_t index = threadIdx.x + blockDim.x * blockIdx.x;
    if(index < state_size){
        // f(x)=e^-x/1+e^-x
        fx[index] = expf(-x[index]) / (1. + expf(-x[index]));
    }
}

使用 threadIdx.x + blockDim.x * blockIdx.x 计算全局一维索引，并通过 if(index < state_size) 防止越界——相比文档示例中的二维索引写法，这种方式更通用、更安全。

• host 侧根据元素总数自动计算 block 数：

__host__ at::Tensor sigmoid_cuda_forward(at::Tensor& x) {
    auto fx = x.clone();
    const int state_size = x.numel();
    const int nblocks = (state_size + THREADS - 1) / THREADS;
    AT_DISPATCH_FLOATING_TYPES(x.type(), "sigmoid_forward_cuda", ([&] {
        sigmoid_cuda_forward_kernel<scalar_t><<<nblocks, THREADS>>>(
            x.data<scalar_t>(), fx.data<scalar_t>(), state_size);
  }));
  return fx;
}

用 x.numel() 拿到元素总数，以 (state_size + THREADS - 1) / THREADS 向上取整计算 block 数，并将每个 block 设为 1024 个线程，是更贴近实际工程的标准写法。

• C++ 绑定层的输入检查宏（见 Chapter8/CUDAExtension/src/MySigmoidCUDA.cpp）：

#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x "must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x "must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

在转发给 CUDA 函数前先断言输入必须是 CUDA 张量且内存连续，避免隐式拷贝与 kernel 数据错位。

3.5 CUDA 扩展的编译安装

与 C++ 扩展不同，CUDA 扩展除了使用 CppExtension 以外，还会同时使用 CUDAExtension。构建时也可以采用 JIT 实时扩展，但这里仍然使用 setuptools 方式：

# ./setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension

setup(
    name='linear_cuda', # 模块名称，用于import调用
    ext_modules=[
        CUDAExtension('linear_cuda',
        [
            './src/linear_cuda_kernel.cu',
            './src/linear_cuda.cpp',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

此时整个工程的目录结构为：

|__ CUDAExtension
    |—— src
        |—— linear_cuda_kernel.cu
        |__ linear_cuda.cpp
    |__ setup.py

在 setup.py 所在目录执行 python setup.py install 完成构建，模块会安装到 Python 的 site-packages 中。看到类似如下输出即表示构建成功：

... ...
Installed /home/admin/anaconda3/lib/python3.7/site-packages/linear_cuda-0.0.0-py3.7-linux-x86_64.egg
Processing dependencies for linear-cuda==0.0.0
Finished processing dependencies for linear-cuda==0.0.0

本仓库实际的 CUDA 扩展构建脚本见 Chapter8/CUDAExtension/setup.py，模块名为 mysigmoid2，其 sources 列表同样同时包含 .cu 内核文件与 C++ 绑定文件。

3.6 验证 CUDA 扩展的正确性

# ./test.py
import torch
from torch.autograd import Function
import linear_cuda # 导入模块

class Linear_CUDA(Function):
    @staticmethod
    def forward(ctx, x, b):
        t_x, t_y = linear_cuda.forward(x, b)
        vars = [t_x]
        ctx.save_for_backward(*vars)
        return t_y

    @staticmethod
    def backward(ctx, grad_y):
        grad_x, grad_b = linear_cuda.backward(*ctx.saved_tensors, grad_y)
        return grad_x, grad_b


class Test(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, b):
        return Linear_CUDA.apply(x, b)

def main():
    x1 = torch.randn(500, 500).cuda().requires_grad_()
    b1 = torch.randn(500, 500).cuda().requires_grad_()
    model = Test().cuda()
    y1 = model(x1, b1)
    y1.sum().backward()

    x2 = x1.clone().detach().cuda().requires_grad_()
    b2 = b1.clone().detach().cuda().requires_grad_()
    y2 = x2 ** 2 + b2
    y2.sum().backward()

    assert y1.equal(y2)
    assert x1.grad.equal(x2.grad)
    assert b1.grad.equal(b2.grad)

if __name__ == '__main__':
    main()

注意与 C++ 版本的区别：CUDA 版本的输入输出都是 .cuda() 张量，且模型也需 .cuda() 上卡。运行测试没有异常，即表示自定义 CUDA 扩展与 PyTorch 原生反向传播结果一致。

仓库配套测试 Chapter8/CUDAExtension/test.py 在断言时使用 fx1.data.cpu().allclose(fx2) 把 GPU 结果拷回 CPU 再比较，并额外内置了 compare_pytorch_with_cuda()，对 (1280, 1280) 的 CUDA 张量分别计时 CUDA 扩展与 PyTorch 向量化的 forward/backward 耗时。

3.7 两种扩展方式总结

至此我们已经完成了 PyTorch 的 C++ 扩展与 CUDA 扩展。二者在构建上异曲同工：

• 都是将重写好的前向传播函数与反向传播函数绑定到 Python 上；
• 都通过 install 安装成可 import 的包；
• 最后在调用时，利用该包实例化一个自定义的 autograd.Function 对象。

3.8 性能比较

我们测试了一个高维矩阵进行矩阵相乘（逐点运算）的场景：

• 相比直接使用 Python 原生的 For 循环，C++ 扩展省去了大量上层调用，速度有了明显提升；
• CUDA 扩展有效利用了内核多线程并行计算的特性，将时间复杂度高的串行计算优化为速度更快的并行计算，性能在三种方案中最优。

这一点也可以通过仓库 Chapter8/CppExtension/test.py 与 Chapter8/CUDAExtension/test.py 内置的计时函数自行复现验证（两者均对 (1280, 1280) 张量计时）。

四、CUDA / NVIDIA-driver / cuDNN / Python 之间的关系

PyTorch 与 CUDA 更新迭代速度很快，实际使用中常会遇到 CUDA 运行版本、编译版本与框架代码不匹配的问题，本小节系统梳理这几个概念。

4.1 NVIDIA-driver 与 CUDA Toolkit 的版本对应

NVIDIA-driver 是针对 NVIDIA GPU 的驱动程序，安装 PyTorch、CUDA 之前必须保证驱动已装好。NVIDIA-driver 向下兼容：系统驱动版本决定了系统能支持什么版本的 CUDA 与 cudatoolkit。因此驱动版本不需要刻意与 CUDA 版本对齐，保持一个较高的版本即可。部分 CUDA Toolkit 与 NVIDIA-driver 的版本对应关系如下：

CUDA Toolkit	Linux x86_64 Driver Version	Windows x86_64 Driver Version
CUDA 11.1 GA	>=455.23	>=456.38
CUDA 11.0.1 RC	>= 450.36.06	>= 451.22
CUDA 10.1 (10.1.105)	>= 418.39	>= 418.96
CUDA 10.0.130	>= 410.48	>= 411.31
CUDA 9.0 (9.0.76)	>= 384.81	>= 385.54
CUDA 8.0 (8.0.44)	>= 367.48	>= 369.30
CUDA 7.5 (7.5.16)	>= 352.31	>= 353.66

4.2 cudatoolkit 与官方 CUDA Toolkit 的区别

通过 Anaconda 安装 GPU 版 PyTorch 时，安装目录下会自动出现 cudatoolkit。它与官方提供的 CUDA Toolkit 存在差异：

• 官方 CUDA Toolkit 是完整的工具安装包，包含 CUDA-C / CUDA-C++ 编译器以及 nvcc。CUDA 程序包含两种代码：运行在 CPU 上的 host 代码与运行在 GPU 上的 device 代码，nvcc 保证两部分代码的编译结果能在不同机器上运行。Linux 下可用 nvcc --version 或 nvcc -V 查看：

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2017 NVIDIA Corporation
Built on Fri_Nov__3_21:07:56_CDT_2017
Cuda compilation tools, release 9.1, V9.1.85

• CUDA Toolkit 还包含 CUDA 开发所需的编译、调试等组件。但对于 PyTorch 而言这些组件不是那么重要：PyTorch 使用 GPU 时大多数情况下只是调用 CUDA 的动态链接库来支持程序运行，PyTorch 部分与 CUDA 部分的代码是提前编译好的，运行过程不需要重新编译，直接在动态链接库中执行即可。
• Anaconda 安装的 cudatoolkit 里包含应用使用 CUDA 功能所需的动态链接库，此时不需要再安装官方 CUDA Toolkit；
• 但若要给 PyTorch 添加 CUDA 扩展并对编写的 CUDA 程序进行编译，就必须安装官方提供的完整 CUDA Toolkit 工具包。

4.3 nvidia-smi 与两个 CUDA 版本号之谜

在 NVIDIA-driver 支持下，可用 nvidia-smi 查看版本信息，同时管理和监控 GPU 设备：

Tue Dec  8 10:45:19 2020
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 430.26       Driver Version: 430.26       CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 108...  Off  | 00000000:02:00.0 Off |                  N/A |
|  0%   60C    P5    27W / 280W |      0MiB / 11177MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

对比 nvcc --version（CUDA 9.1）与 nvidia-smi（CUDA 10.2），版本号不同但代码仍能正常运行——因为 CUDA 中有两个 API：runtime API 与 driver API，各自都有对应版本号：

• 支持 driver API 的是通过 GPU 驱动安装的，nvidia-smi 就属于这一类；
• 支持 runtime API 的是通过 CUDA Toolkit 安装的。CUDA Toolkit Installer 有时会集成 GPU driver Installer，而 nvcc 随 CUDA Toolkit 一起安装，只反映 CUDA 的 runtime API 版本。

因此这两个版本号不一定一致。CUDA 中几种 API 的关系如图 8-2 所示。

图 8-2 CUDA 中几种 API 之间的关系（来源：Chapter8/img/DRIVER.png）

一个应用只能使用其中一个 API：

• runtime API 拥有高层级的封装，运行时可以编译并将 CUDA 内核链接到可执行文件中，通过隐式初始化、上下文管理和模块管理简化设备代码管理；
• driver API 更接近底层，编程实现更困难，但能查询到更详细的设备信息。

4.4 cuDNN：为深度学习设计的 GPU 加速库

cuDNN 是专门为深度学习设计的软件库，提供了深度神经网络 GPU 加速库，内含大量封装好的计算函数（如卷积、池化等），目标是保障性能、可用性以及更低的内存开销。英伟达的 cuDNN 可以集成到 PyTorch、TensorFlow 等高级机器学习框架中。

cuDNN 通过简单的插入式操作便能使用——只需将 cuDNN 的文件复制到 CUDA 对应的文件夹里，即完成 CUDA 扩展，不会对 CUDA 本身造成其他影响。用户无需浪费时间在 GPU 性能调优上，即可在 GPU 上实现高性能并行计算。

五、本章小结

本章围绕 PyTorch 的自定义算子扩展做了完整的梳理：

• 介绍了 C++ 扩展的完整三步流程：编写 C++ 函数并绑定 pybind11 → setuptools 编译安装 → 用 autograd.Function 封装后放入 nn.Module；
• 介绍了 CUDA 扩展：在 C++ 绑定层之外，编写以 .cu 结尾的 CUDA 内核（掌握 Grid/Block/Thread 三级线程组织、__global__/__device__/__host__ 声明、AT_DISPATCH_FLOATING_TYPES 特化调用），再用 CUDAExtension 编译；
• 通过一个 y = x² + b 的示例演示了两种扩展方式，并对比了 Python 循环、C++ 扩展、CUDA 扩展的性能差异；
• 剖析了 CUDA 运行版本、编译版本与代码不匹配的问题：NVIDIA-driver 向下兼容，读者实际使用时只需保证一个较高的 NVIDIA-driver 版本；区分了 cudatoolkit 与官方 CUDA Toolkit（nvcc）、runtime API 与 driver API，以及 cuDNN 的插入式集成方式。

配套的可运行源码与测试位于 Chapter8/CppExtension（C++ 扩展 Sigmoid 示例）与 Chapter8/CUDAExtension（CUDA 扩展 Sigmoid 示例），读者可结合本章讲解逐步编译、运行与验证。