自定义 C++ 和 CUDA 扩展

本教程自 PyTorch 2.4 起已弃用。有关使用自定义 C++/CUDA 扩展扩展 PyTorch 的最新指南，请参阅 PyTorch 自定义运算符。

PyTorch 提供了大量与神经网络、任意张量代数、数据处理及其他用途相关的操作。然而，您可能仍然会遇到需要更定制化操作的情况。例如，您可能希望使用在论文中发现的新颖激活函数，或者实现您作为研究一部分开发的操作。

在 PyTorch 中集成此类自定义操作的最简单方法是使用 Python 编写，通过扩展 Function 和 Module，如此处所述。这将赋予您自动微分的全部功能（无需编写导数函数）以及 Python 通常的表达能力。然而，有时您的操作可能更适合用 C++ 实现。例如，您的代码可能需要非常快速，因为它在模型中频繁调用，或者即使调用次数很少也非常耗时。另一个可能的原因是它依赖于或与其他 C 或 C++ 库交互。为了解决这些情况，PyTorch 提供了一种非常简单的方式来编写自定义的 C++ 扩展。

C++ 扩展是我们开发的一种机制，允许用户（您）创建在外部定义的 PyTorch 操作符，即与 PyTorch 后端分离。这种方法与原生 PyTorch 操作的实现方式不同。C++ 扩展旨在减少将操作与 PyTorch 后端集成时的大量样板代码，同时为您的基于 PyTorch 的项目提供高度的灵活性。尽管如此，一旦您将操作定义为 C++ 扩展，将其转换为原生 PyTorch 函数主要是代码组织的问题，如果您决定将您的操作贡献到上游，可以在事后处理。

动机与示例

本文的剩余部分将通过一个实际的示例来介绍如何编写和使用C++（以及CUDA）扩展。如果你正被追赶，或者有人威胁你如果今天不完成这个操作就会解雇你，你可以跳过本节，直接阅读下一节中的实现细节。

假设你发明了一种新型的循环单元，发现它比现有的技术具有更优越的特性。这个循环单元类似于LSTM，但不同之处在于它没有遗忘门，并且使用指数线性单元（ELU）作为其内部激活函数。由于这个单元永远不会遗忘，我们将其称为LLTM，即长长期记忆单元。

LLTM与普通的LSTM有两个显著的不同之处，以至于我们无法通过配置PyTorch的LSTMCell来满足我们的需求，因此我们必须创建一个自定义的单元。实现这一目标的第一种也是最简单的方法——可能在所有情况下都是一个良好的第一步——是在纯PyTorch中使用Python实现我们所需的功能。为此，我们需要继承torch.nn.Module并实现LLTM的前向传递。这将如下所示：

classLLTM(torch.nn.Module):
    def__init__(self, input_features, state_size):
        super(LLTM, self).__init__()
        self.input_features = input_features
        self.state_size = state_size
        # 3 * state_size for input gate, output gate and candidate cell gate.
        # input_features + state_size because we will multiply with [input, h].
        self.weights = torch.nn.Parameter(
            torch.empty(3 * state_size, input_features + state_size))
        self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
        self.reset_parameters()

    defreset_parameters(self):
        stdv = 1.0 / math.sqrt(self.state_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, +stdv)

    defforward(self, input, state):
        old_h, old_cell = state
        X = torch.cat([old_h, input], dim=1)

        # Compute the input, output and candidate cell gates with one MM.
        gate_weights = F.linear(X, self.weights, self.bias)
        # Split the combined gate weight matrix into its components.
        gates = gate_weights.chunk(3, dim=1)

        input_gate = torch.sigmoid(gates[0])
        output_gate = torch.sigmoid(gates[1])
        # Here we use an ELU instead of the usual tanh.
        candidate_cell = F.elu(gates[2])

        # Compute the new cell state.
        new_cell = old_cell + candidate_cell * input_gate
        # Compute the new hidden state and output.
        new_h = torch.tanh(new_cell) * output_gate

        return new_h, new_cell

我们可以按预期使用它：

importtorch

X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)

rnn = LLTM(input_features, state_size)

new_h, new_C = rnn(X, (h, C))

当然，如果可能且合理的话，您应该使用这种方法来扩展 PyTorch。由于 PyTorch 为其 CPU 和 GPU 操作提供了高度优化的实现，这些实现由 NVIDIA cuDNN、Intel MKL 或 NNPACK 等库提供支持，因此像上面这样的 PyTorch 代码通常已经足够快。然而，我们也可以看到，在某些情况下，仍然存在进一步性能提升的空间。最明显的原因是 PyTorch 并不了解您正在实现的算法。它只知道您用于构建算法的各个操作。因此，PyTorch 必须逐个执行您的操作。由于每个单独的操作实现（或内核）调用可能涉及 CUDA 内核的启动，这会产生一定的开销，在多次函数调用中，这种开销可能会变得显著。此外，运行我们代码的 Python 解释器本身也可能会拖慢程序的执行速度。

因此，加速的明确方法是将部分代码用 C++（或 CUDA）重写，并将特定的操作组进行融合。融合意味着将多个函数的实现合并为一个单一的函数，这样可以减少内核启动次数，并且由于对数据全局流的可见性增加，还能进行其他优化。

让我们看看如何使用 C++ 扩展来实现 LLTM 的融合版本。我们将首先使用 ATen 库以纯 C++ 编写它，ATen 是 PyTorch 后端的核心库，它使我们能够轻松地将 Python 代码转换为 C++。然后，我们将模型的部分代码迁移到 CUDA 内核中，以利用 GPU 提供的强大并行计算能力，从而进一步提升性能。

编写 C++ 扩展

C++ 扩展有两种形式：可以通过 setuptools 提前构建，也可以通过 torch.utils.cpp_extension.load() 即时加载。我们将从第一种方法开始，稍后再讨论后者。

使用 `setuptools` 构建

对于“提前编译”的方式，我们通过编写一个使用 setuptools 来编译 C++ 代码的 setup.py 脚本来构建我们的 C++ 扩展。对于 LLTM，它的实现看起来非常简单：

fromsetuptoolsimport setup, Extension
fromtorch.utilsimport cpp_extension

setup(name='lltm_cpp',
      ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['lltm.cpp'])],
      cmdclass={'build_ext': cpp_extension.BuildExtension})

在这段代码中，CppExtension 是 setuptools.Extension 的一个便捷封装，它传递了正确的包含路径并将扩展的语言设置为 C++。等效的原生 setuptools 代码将简单地是：

Extension(
   name='lltm_cpp',
   sources=['lltm.cpp'],
   include_dirs=cpp_extension.include_paths(),
   language='c++')

BuildExtension 执行了一系列必需的配置步骤和检查，并且在混合 C++/CUDA 扩展的情况下管理混合编译。这就是我们现在关于构建 C++ 扩展所需要了解的全部内容！现在让我们来看看我们 C++ 扩展的实现，这些实现位于 lltm.cpp 中。

编写 C++ 操作

让我们开始用 C++ 实现 LLTM 吧！在反向传播中我们需要的一个函数是 sigmoid 的导数。这段代码足够简短，可以用来讨论编写 C++ 扩展时可用的整体环境：

#include<torch/extension.h>

#include<iostream>

torch::Tensord_sigmoid(torch::Tensorz){
autos=torch::sigmoid(z);
return(1-s)*s;
}

<torch/extension.h> 是一个一站式头文件，包含了编写 PyTorch C++ 扩展所需的所有必要组件。它包括：

ATen 库，这是我们用于张量计算的主要 API，
pybind11，这是我们为 C++ 代码创建 Python 绑定的工具，
管理 ATen 和 pybind11 之间交互细节的头文件。

d_sigmoid() 的实现展示了如何使用 ATen API。PyTorch 的张量和变量接口是从 ATen 库自动生成的，因此我们可以几乎将 Python 实现 1:1 地翻译为 C++。我们所有计算的主要数据类型将是 torch::Tensor。它的完整 API 可以在这里查看 https://pytorch.org/cppdocs/api/classat_1_1_tensor.html。另外需要注意的是，我们可以包含 <iostream> 或 任何其他 C 或 C++ 头文件 —— 我们可以充分利用 C++11 的全部功能。

请注意，CUDA-11.5 的 nvcc 在解析 torch/extension.h 时会在 Windows 上遇到内部编译器错误。要解决此问题，请将 Python 绑定逻辑移至纯 C++ 文件中。示例用法：

#include<ATen/ATen.h>
at::TensorSigmoidAlphaBlendForwardCuda(....)

无需：

#include<torch/extension.h>
torch::TensorSigmoidAlphaBlendForwardCuda(...)

目前开放的 nvcc 问题在这里。完整的解决方案代码示例在这里。

前向传播

接下来，我们可以将整个前向传播过程移植到 C++ 中：

#include<vector>

std::vector<at::Tensor>lltm_forward(
torch::Tensorinput,
torch::Tensorweights,
torch::Tensorbias,
torch::Tensorold_h,
torch::Tensorold_cell){
autoX=torch::cat({old_h,input},/*dim=*/1);

autogate_weights=torch::addmm(bias,X,weights.transpose(0,1));
autogates=gate_weights.chunk(3,/*dim=*/1);

autoinput_gate=torch::sigmoid(gates[0]);
autooutput_gate=torch::sigmoid(gates[1]);
autocandidate_cell=torch::elu(gates[2],/*alpha=*/1.0);

autonew_cell=old_cell+candidate_cell*input_gate;
autonew_h=torch::tanh(new_cell)*output_gate;

return{new_h,
new_cell,
input_gate,
output_gate,
candidate_cell,
X,
gate_weights};
}

反向传播

C++扩展API目前并未提供自动生成反向函数的方法。因此，我们还需要实现LLTM的反向传播，以计算损失相对于前向传播输入的导数。最终，我们将把前向和反向函数放入torch.autograd.Function中，以创建一个良好的Python绑定。反向函数稍微复杂一些，因此我们不会深入探讨代码（如果你对此感兴趣，Alex Graves的论文是获取更多信息的好资源）：

// tanh'(z) = 1 - tanh^2(z)
torch::Tensord_tanh(torch::Tensorz){
return1-z.tanh().pow(2);
}

// elu'(z) = relu'(z) + { alpha * exp(z) if (alpha * (exp(z) - 1)) < 0, else 0}
torch::Tensord_elu(torch::Tensorz,torch::Scalaralpha=1.0){
autoe=z.exp();
automask=(alpha*(e-1))<0;
return(z>0).type_as(z)+mask.type_as(z)*(alpha*e);
}

std::vector<torch::Tensor>lltm_backward(
torch::Tensorgrad_h,
torch::Tensorgrad_cell,
torch::Tensornew_cell,
torch::Tensorinput_gate,
torch::Tensoroutput_gate,
torch::Tensorcandidate_cell,
torch::TensorX,
torch::Tensorgate_weights,
torch::Tensorweights){
autod_output_gate=torch::tanh(new_cell)*grad_h;
autod_tanh_new_cell=output_gate*grad_h;
autod_new_cell=d_tanh(new_cell)*d_tanh_new_cell+grad_cell;

autod_old_cell=d_new_cell;
autod_candidate_cell=input_gate*d_new_cell;
autod_input_gate=candidate_cell*d_new_cell;

autogates=gate_weights.chunk(3,/*dim=*/1);
d_input_gate*=d_sigmoid(gates[0]);
d_output_gate*=d_sigmoid(gates[1]);
d_candidate_cell*=d_elu(gates[2]);

autod_gates=
torch::cat({d_input_gate,d_output_gate,d_candidate_cell},/*dim=*/1);

autod_weights=d_gates.t().mm(X);
autod_bias=d_gates.sum(/*dim=*/0,/*keepdim=*/true);

autod_X=d_gates.mm(weights);
constautostate_size=grad_h.size(1);
autod_old_h=d_X.slice(/*dim=*/1,0,state_size);
autod_input=d_X.slice(/*dim=*/1,state_size);

return{d_old_h,d_input,d_weights,d_bias,d_old_cell};
}

绑定到 Python

一旦您用 C++ 和 ATen 编写了操作，您可以使用 pybind11 以非常简洁的方式将 C++ 函数或类绑定到 Python 中。关于 PyTorch C++ 扩展的这部分问题或疑问，大部分都可以在 pybind11 文档中找到解答。

对于我们的扩展，必要的绑定代码仅需四行：

PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){
m.def("forward",&lltm_forward,"LLTM forward");
m.def("backward",&lltm_backward,"LLTM backward");
}

这里需要注意的一点是宏 TORCH_EXTENSION_NAME。torch 扩展构建会将其定义为您在 setup.py 脚本中为扩展指定的名称。在本例中，TORCH_EXTENSION_NAME 的值将是“lltm_cpp”。这样做是为了避免在两个地方（构建脚本和 C++ 代码）维护扩展名称，因为两者不匹配可能导致难以追踪的严重问题。

使用您的扩展

我们现在准备在 PyTorch 中导入我们的扩展。此时，您的目录结构可能如下所示：

pytorch/
  lltm-extension/
    lltm.cpp
    setup.py

现在，运行 python setup.py install 来构建并安装您的扩展。结果应该类似于以下内容：

running install
running bdist_egg
running egg_info
creating lltm_cpp.egg-info
writing lltm_cpp.egg-info/PKG-INFO
writing dependency_links to lltm_cpp.egg-info/dependency_links.txt
writing top-level names to lltm_cpp.egg-info/top_level.txt
writing manifest file 'lltm_cpp.egg-info/SOURCES.txt'
reading manifest file 'lltm_cpp.egg-info/SOURCES.txt'
writing manifest file 'lltm_cpp.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_ext
building 'lltm_cpp' extension
creating build
creating build/temp.linux-x86_64-3.7
gcc -pthread -B ~/local/miniconda/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I~/local/miniconda/lib/python3.7/site-packages/torch/include -I~/local/miniconda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I~/local/miniconda/lib/python3.7/site-packages/torch/include/TH -I~/local/miniconda/lib/python3.7/site-packages/torch/include/THC -I~/local/miniconda/include/python3.7m -c lltm.cpp -o build/temp.linux-x86_64-3.7/lltm.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=lltm_cpp -D_GLIBCXX_USE_CXX11_ABI=1 -std=c++11
cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++
creating build/lib.linux-x86_64-3.7
g++ -pthread -shared -B ~/local/miniconda/compiler_compat -L~/local/miniconda/lib -Wl,-rpath=~/local/miniconda/lib -Wl,--no-as-needed -Wl,--sysroot=/ build/temp.linux-x86_64-3.7/lltm.o -o build/lib.linux-x86_64-3.7/lltm_cpp.cpython-37m-x86_64-linux-gnu.so
creating build/bdist.linux-x86_64
creating build/bdist.linux-x86_64/egg
copying build/lib.linux-x86_64-3.7/lltm_cpp.cpython-37m-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/egg
creating stub loader for lltm_cpp.cpython-37m-x86_64-linux-gnu.so
byte-compiling build/bdist.linux-x86_64/egg/lltm_cpp.py to lltm_cpp.cpython-37.pyc
creating build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying lltm_cpp.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
writing build/bdist.linux-x86_64/egg/EGG-INFO/native_libs.txt
zip_safe flag not set; analyzing archive contents...
__pycache__.lltm_cpp.cpython-37: module references __file__
creating 'dist/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg' and adding 'build/bdist.linux-x86_64/egg' to it
removing 'build/bdist.linux-x86_64/egg' (and everything under it)
Processing lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
removing '~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg' (and everything under it)
creating ~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
Extracting lltm_cpp-0.0.0-py3.7-linux-x86_64.egg to ~/local/miniconda/lib/python3.7/site-packages
lltm-cpp 0.0.0 is already the active version in easy-install.pth

Installed ~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
Processing dependencies for lltm-cpp==0.0.0
Finished processing dependencies for lltm-cpp==0.0.0

关于编译器的一个小注意事项：由于ABI版本兼容性问题，您用来构建C++扩展的编译器必须与构建PyTorch的编译器ABI兼容。实际上，这意味着在Linux上您必须使用GCC 4.9及以上版本。对于Ubuntu 16.04及更新版本的Linux发行版，这通常已经是默认的编译器。在MacOS上，您必须使用clang（它没有ABI版本兼容性问题）。在最坏的情况下，您可以使用自己的编译器从源代码构建PyTorch，然后用相同的编译器构建扩展。

当您的扩展构建完成后，您可以在 Python 中轻松导入它，使用您在 setup.py 脚本中指定的名称。只需确保首先 import torch，因为这将解析动态链接器必须看到的一些符号：

In [1]: importtorch
In [2]: importlltm_cpp
In [3]: lltm_cpp.forward
Out[3]: <function lltm.PyCapsule.forward>

如果我们对函数或模块调用 help()，可以看到它的签名与我们的 C++ 代码相匹配：

In[4] help(lltm_cpp.forward)
forward(...) method of builtins.PyCapsule instance
    forward(arg0: torch::Tensor, arg1: torch::Tensor, arg2: torch::Tensor, arg3: torch::Tensor, arg4: torch::Tensor) -> List[torch::Tensor]

    LLTM forward

由于我们现在能够从 Python 调用 C++ 函数，我们可以使用 torch.autograd.Function 和 torch.nn.Module 对它们进行封装，使它们成为 PyTorch 的一等公民：

importmath
importtorch

# Our module!
importlltm_cpp

classLLTMFunction(torch.autograd.Function):
    @staticmethod
    defforward(ctx, input, weights, bias, old_h, old_cell):
        outputs = lltm_cpp.forward(input, weights, bias, old_h, old_cell)
        new_h, new_cell = outputs[:2]
        variables = outputs[1:] + [weights]
        ctx.save_for_backward(*variables)

        return new_h, new_cell

    @staticmethod
    defbackward(ctx, grad_h, grad_cell):
        outputs = lltm_cpp.backward(
            grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_tensors)
        d_old_h, d_input, d_weights, d_bias, d_old_cell = outputs
        return d_input, d_weights, d_bias, d_old_h, d_old_cell


classLLTM(torch.nn.Module):
    def__init__(self, input_features, state_size):
        super(LLTM, self).__init__()
        self.input_features = input_features
        self.state_size = state_size
        self.weights = torch.nn.Parameter(
            torch.empty(3 * state_size, input_features + state_size))
        self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
        self.reset_parameters()

    defreset_parameters(self):
        stdv = 1.0 / math.sqrt(self.state_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, +stdv)

    defforward(self, input, state):
        return LLTMFunction.apply(input, self.weights, self.bias, *state)

性能对比

现在我们已经能够从 PyTorch 中使用和调用我们的 C++ 代码了，我们可以运行一个小型基准测试，看看通过用 C++ 重写我们的操作获得了多少性能提升。我们将多次运行 LLTM 的前向和后向传播，并测量其耗时：

importtime

importtorch

batch_size = 16
input_features = 32
state_size = 128

X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)

rnn = LLTM(input_features, state_size)

forward = 0
backward = 0
for _ in range(100000):
    start = time.time()
    new_h, new_C = rnn(X, (h, C))
    forward += time.time() - start

    start = time.time()
    (new_h.sum() + new_C.sum()).backward()
    backward += time.time() - start

print('Forward: {:.3f} s | Backward {:.3f} s'.format(forward, backward))

如果我们用本文开头用纯 Python 编写的原始 LLTM 来运行这段代码，将得到以下数值（在我的机器上）：

Forward: 506.480 us | Backward 444.694 us

并使用我们的新 C++ 版本：

Forward: 349.335 us | Backward 443.523 us

我们已经可以看到前向函数的速度显著提升（超过30%）。对于反向函数，虽然速度有所提升，但并不显著。上面我编写的反向传递代码并没有特别优化，肯定还有改进的空间。此外，PyTorch的自动微分引擎可以自动并行化计算图，可能整体上使用了更高效的操作流程，并且也是用C++实现的，因此它的速度预计会很快。不过，这已经是一个很好的开始了。

GPU 设备上的性能

PyTorch 的 ATen 后端有一个非常棒的特性，那就是它抽象了您正在运行的计算设备。这意味着我们为 CPU 编写的代码同样可以在 GPU 上运行，并且各个操作会相应地分派到 GPU 优化的实现上。对于某些操作，如矩阵乘法（如 mm 或 addmm），这带来了巨大的性能提升。让我们看看在使用 CUDA 张量运行我们的 C++ 代码时能获得多少性能提升。我们不需要对实现进行任何更改，只需将我们的张量放入 GPU 内存中即可。可以通过在创建时添加 device=cuda_device 参数，或者在创建后使用 .to(cuda_device) 来实现：

importtorch

assert torch.cuda.is_available()
cuda_device = torch.device("cuda")  # device object representing GPU

batch_size = 16
input_features = 32
state_size = 128

# Note the device=cuda_device arguments here
X = torch.randn(batch_size, input_features, device=cuda_device)
h = torch.randn(batch_size, state_size, device=cuda_device)
C = torch.randn(batch_size, state_size, device=cuda_device)

rnn = LLTM(input_features, state_size).to(cuda_device)

forward = 0
backward = 0
for _ in range(100000):
    start = time.time()
    new_h, new_C = rnn(X, (h, C))
    torch.cuda.synchronize()
    forward += time.time() - start

    start = time.time()
    (new_h.sum() + new_C.sum()).backward()
    torch.cuda.synchronize()
    backward += time.time() - start

print('Forward: {:.3f} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))

再次将我们的纯 PyTorch 代码与 C++ 版本进行比较，现在两者都在 CUDA 设备上运行，我们再次看到了性能提升。对于 Python/PyTorch：

Forward: 187.719 us | Backward 410.815 us

以及 C++/ATen：

Forward: 149.802 us | Backward 393.458 us

与非CUDA代码相比，这是一个显著的加速。然而，通过编写自定义的CUDA内核，我们可以进一步提升C++代码的性能，这一点我们很快就会深入探讨。在此之前，让我们先讨论另一种构建C++扩展的方法。

JIT 编译扩展

之前，我提到过构建 C++ 扩展的两种方法：使用 setuptools 或即时编译（JIT）。在介绍了前一种方法之后，让我们详细讨论一下后者。JIT 编译机制为您提供了一种通过调用 PyTorch API 中的一个简单函数 torch.utils.cpp_extension.load() 来动态编译和加载扩展的方式。对于 LLTM 来说，这将非常简单，如下所示：

fromtorch.utils.cpp_extensionimport load

lltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

在这里，我们向函数提供了与 setuptools 相同的信息。在后台，这将执行以下操作：

创建一个临时目录 /tmp/torch_extensions/lltm，
在该临时目录中生成一个 Ninja 构建文件，
将您的源文件编译成一个共享库，
将此共享库作为 Python 模块导入。

实际上，如果您将 verbose=True 传递给 cpp_extension.load()，您将会收到有关过程的详细信息：

Using /tmp/torch_extensions as PyTorch extensions root...
Emitting ninja build file /tmp/torch_extensions/lltm_cpp/build.ninja...
Building extension module lltm_cpp...
Loading extension module lltm_cpp...

生成的 Python 模块将与 setuptools 生成的完全相同，但无需维护单独的 setup.py 构建文件。如果您的设置更为复杂，确实需要 setuptools 的全部功能，您可以编写自己的 setup.py——但在许多情况下，这种即时编译（JIT）技术就足够了。第一次运行这行代码时，由于扩展程序在后台编译，会花费一些时间。由于我们使用 Ninja 构建系统来构建您的源代码，重新编译是增量的，因此当您第二次运行 Python 模块时，重新加载扩展的速度很快，并且如果您没有更改扩展的源文件，开销也很低。

编写混合 C++/CUDA 扩展

为了将我们的实现提升到一个新的水平，我们可以使用自定义的 CUDA 内核手动编写部分前向传播和反向传播过程。对于 LLTM（Long Short-Term Memory）来说，这种方法尤其有效，因为其中包含大量的按顺序执行的点操作，这些操作都可以融合并在单个 CUDA 内核中并行化。接下来，我们将探讨如何编写这样的 CUDA 内核，并通过这种扩展机制将其与 PyTorch 集成。

编写 CUDA 扩展的一般策略是首先编写一个 C++ 文件，该文件定义了将从 Python 调用的函数，并使用 pybind11 将这些函数绑定到 Python。此外，该文件还会声明在 CUDA (.cu) 文件中定义的函数。C++ 函数将执行一些检查，并最终将其调用转发给 CUDA 函数。在 CUDA 文件中，我们编写实际的 CUDA 内核。然后，cpp_extension 包将负责使用 C++ 编译器（如 gcc）编译 C++ 源文件，并使用 NVIDIA 的 nvcc 编译器编译 CUDA 源文件。这确保了每个编译器都能处理其最擅长的文件类型。最终，它们将被链接到一个共享库中，供我们从 Python 代码中调用。

我们将从C++文件开始，例如，我们将其命名为lltm_cuda.cpp：

#include<torch/extension.h>

#include<vector>

// CUDA forward declarations

std::vector<torch::Tensor>lltm_cuda_forward(
torch::Tensorinput,
torch::Tensorweights,
torch::Tensorbias,
torch::Tensorold_h,
torch::Tensorold_cell);

std::vector<torch::Tensor>lltm_cuda_backward(
torch::Tensorgrad_h,
torch::Tensorgrad_cell,
torch::Tensornew_cell,
torch::Tensorinput_gate,
torch::Tensoroutput_gate,
torch::Tensorcandidate_cell,
torch::TensorX,
torch::Tensorgate_weights,
torch::Tensorweights);

// C++ interface

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

std::vector<torch::Tensor>lltm_forward(
torch::Tensorinput,
torch::Tensorweights,
torch::Tensorbias,
torch::Tensorold_h,
torch::Tensorold_cell){
CHECK_INPUT(input);
CHECK_INPUT(weights);
CHECK_INPUT(bias);
CHECK_INPUT(old_h);
CHECK_INPUT(old_cell);

returnlltm_cuda_forward(input,weights,bias,old_h,old_cell);
}

std::vector<torch::Tensor>lltm_backward(
torch::Tensorgrad_h,
torch::Tensorgrad_cell,
torch::Tensornew_cell,
torch::Tensorinput_gate,
torch::Tensoroutput_gate,
torch::Tensorcandidate_cell,
torch::TensorX,
torch::Tensorgate_weights,
torch::Tensorweights){
CHECK_INPUT(grad_h);
CHECK_INPUT(grad_cell);
CHECK_INPUT(input_gate);
CHECK_INPUT(output_gate);
CHECK_INPUT(candidate_cell);
CHECK_INPUT(X);
CHECK_INPUT(gate_weights);
CHECK_INPUT(weights);

returnlltm_cuda_backward(
grad_h,
grad_cell,
new_cell,
input_gate,
output_gate,
candidate_cell,
X,
gate_weights,
weights);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){
m.def("forward",&lltm_forward,"LLTM forward (CUDA)");
m.def("backward",&lltm_backward,"LLTM backward (CUDA)");
}

如您所见，这部分代码主要是样板代码、检查操作以及转发到我们将在 CUDA 文件中定义的函数。我们会将此文件命名为 lltm_cuda_kernel.cu（注意 .cu 扩展名！）。NVCC 能够较好地编译 C++11，因此我们仍然可以使用 ATen 和 C++ 标准库（但不能使用 torch.h）。需要注意的是，setuptools 无法处理同名但扩展名不同的文件，因此如果您使用 setup.py 方法而非 JIT 方法，必须为 CUDA 文件起一个与 C++ 文件不同的名称（对于 JIT 方法，lltm.cpp 和 lltm.cu 可以正常工作）。让我们简单看一下这个文件的内容：

#include<torch/extension.h>

#include<cuda.h>
#include<cuda_runtime.h>

#include<vector>

template<typenamescalar_t>
__device____forceinline__scalar_tsigmoid(scalar_tz){
return1.0/(1.0+exp(-z));
}

在这里，我们可以看到我刚刚描述的头部信息，以及我们正在使用 CUDA 特有的声明，例如 __device__ 和 __forceinline__，以及像 exp 这样的函数。接下来，我们将继续介绍一些我们需要的辅助函数：

template<typenamescalar_t>
__device____forceinline__scalar_td_sigmoid(scalar_tz){
constautos=sigmoid(z);
return(1.0-s)*s;
}

template<typenamescalar_t>
__device____forceinline__scalar_td_tanh(scalar_tz){
constautot=tanh(z);
return1-(t*t);
}

template<typenamescalar_t>
__device____forceinline__scalar_telu(scalar_tz,scalar_talpha=1.0){
returnfmax(0.0,z)+fmin(0.0,alpha*(exp(z)-1.0));
}

template<typenamescalar_t>
__device____forceinline__scalar_td_elu(scalar_tz,scalar_talpha=1.0){
constautoe=exp(z);
constautod_relu=z<0.0?0.0:1.0;
returnd_relu+(((alpha*(e-1.0))<0.0)?(alpha*e):0.0);
}

为了实现一个功能，我们需要再次准备两样东西：一个函数，用于执行我们不希望手动编写的操作并调用 CUDA 内核，以及实际的 CUDA 内核，用于加速我们想要优化的部分。对于前向传播，第一个函数应该如下所示：

std::vector<torch::Tensor>lltm_cuda_forward(
torch::Tensorinput,
torch::Tensorweights,
torch::Tensorbias,
torch::Tensorold_h,
torch::Tensorold_cell){
autoX=torch::cat({old_h,input},/*dim=*/1);
autogates=torch::addmm(bias,X,weights.transpose(0,1));

constautobatch_size=old_cell.size(0);
constautostate_size=old_cell.size(1);

autonew_h=torch::zeros_like(old_cell);
autonew_cell=torch::zeros_like(old_cell);
autoinput_gate=torch::zeros_like(old_cell);
autooutput_gate=torch::zeros_like(old_cell);
autocandidate_cell=torch::zeros_like(old_cell);

constintthreads=1024;
constdim3blocks((state_size+threads-1)/threads,batch_size);

AT_DISPATCH_FLOATING_TYPES(gates.type(),"lltm_forward_cuda",([&]{
lltm_cuda_forward_kernel<scalar_t><<<blocks,threads>>>(
gates.data<scalar_t>(),
old_cell.data<scalar_t>(),
new_h.data<scalar_t>(),
new_cell.data<scalar_t>(),
input_gate.data<scalar_t>(),
output_gate.data<scalar_t>(),
candidate_cell.data<scalar_t>(),
state_size);
}));

return{new_h,new_cell,input_gate,output_gate,candidate_cell,X,gates};
}

这里的关键点在于 AT_DISPATCH_FLOATING_TYPES 宏和内核启动（由 <<<...>>> 表示）。虽然 ATen 抽象了我们处理的张量的设备和数据类型，但在运行时，张量仍然由具体设备上的具体类型的内存支持。因此，我们需要一种在运行时确定张量类型的方法，然后有选择地调用具有相应正确类型签名的函数。如果手动完成，这在概念上会类似于以下内容：

switch(tensor.type().scalarType()){
casetorch::ScalarType::Double:
returnfunction<double>(tensor.data<double>());
casetorch::ScalarType::Float:
returnfunction<float>(tensor.data<float>());
...
}

AT_DISPATCH_FLOATING_TYPES 的目的是为我们处理这种分发。它接收一个类型（在我们的例子中是 gates.type()）、一个名称（用于错误消息）和一个 lambda 函数。在这个 lambda 函数内部，类型别名 scalar_t 是可用的，并且被定义为在该上下文中张量在运行时的实际类型。因此，如果我们有一个模板函数（我们的 CUDA 内核将是这样的函数），我们可以使用这个 scalar_t 别名来实例化它，并且会调用正确的函数。在这种情况下，我们还希望将张量的数据指针作为 scalar_t 类型的指针来获取。如果你想分发所有类型而不仅仅是浮点类型（Float 和 Double），你可以使用 AT_DISPATCH_ALL_TYPES。

请注意，我们使用普通的 ATen 执行了一些操作。这些操作仍将在 GPU 上运行，但使用的是 ATen 的默认实现。这是有意义的，因为 ATen 会使用高度优化的例程来处理诸如矩阵乘法（例如 addmm）或卷积等操作，这些操作如果由我们自己实现和改进将会非常困难。

至于内核启动本身，我们在这里指定每个 CUDA 块将有 1024 个线程，并且整个 GPU 网格被分割为尽可能多的 1 x 1024 线程块，以确保我们的矩阵中的每个元素都由一个线程处理。例如，如果我们的状态大小为 2048，批量大小为 4，那么我们将启动总共 4 x 2 = 8 个块，每个块包含 1024 个线程。如果您从未听说过 CUDA 的“块”或“网格”，关于 CUDA 的入门阅读可能会有所帮助。

实际的 CUDA 内核相当简单（如果您之前编写过 GPU 程序的话）：

template<typenamescalar_t>
__global__voidlltm_cuda_forward_kernel(
constscalar_t*__restrict__gates,
constscalar_t*__restrict__old_cell,
scalar_t*__restrict__new_h,
scalar_t*__restrict__new_cell,
scalar_t*__restrict__input_gate,
scalar_t*__restrict__output_gate,
scalar_t*__restrict__candidate_cell,
size_tstate_size){
constintcolumn=blockIdx.x*blockDim.x+threadIdx.x;
constintindex=blockIdx.y*state_size+column;
constintgates_row=blockIdx.y*(state_size*3);
if(column<state_size){
input_gate[index]=sigmoid(gates[gates_row+column]);
output_gate[index]=sigmoid(gates[gates_row+state_size+column]);
candidate_cell[index]=elu(gates[gates_row+2*state_size+column]);
new_cell[index]=
old_cell[index]+candidate_cell[index]*input_gate[index];
new_h[index]=tanh(new_cell[index])*output_gate[index];
}
}

这里主要有趣的是，我们能够完全并行地为门矩阵中的每个单独组件计算所有这些逐点操作。如果你想象必须使用一个包含数百万个元素的巨型for循环串行执行这些操作，你就会明白为什么这种方法会快得多。

使用访问器

您可以在 CUDA 内核中看到，我们直接使用正确类型的指针进行操作。事实上，在 CUDA 内核中直接使用高层级的类型无关张量会非常低效。

然而，这以牺牲易用性和可读性为代价，尤其是对于高维数据。在我们的示例中，我们知道连续的 gates 张量有 3 个维度：

批次，大小为 batch_size，步幅为 3*state_size
行，大小为 3，步幅为 state_size
索引，大小为 state_size，步幅为 1

那么我们如何在核函数中访问元素 gates[n][row][column] 呢？实际上，您可以通过一些简单的算术运算，利用步幅来访问该元素。

gates.data<scalar_t>()[n*3*state_size+row*state_size+column]

除了冗长之外，这种表达式还需要显式知道步幅（stride），因此需要将其作为参数传递给内核函数。您可以看到，在需要接受多个具有不同大小的张量的内核函数中，最终会得到一个非常长的参数列表。

幸运的是，ATen 提供了访问器（accessors），这些访问器通过一次动态检查来确保张量的类型和维度数量正确。访问器随后会公开一个 API，用于高效地访问张量元素，而无需将其转换为单个指针：

torch::Tensorfoo=torch::rand({12,12});

// assert foo is 2-dimensional and holds floats.
autofoo_a=foo.accessor<float,2>();
floattrace=0;

for(inti=0;i<foo_a.size(0);i++){
// use the accessor foo_a to get tensor data.
trace+=foo_a[i][i];
}

访问器对象具有相对高级的接口，包含.size()和.stride()方法以及多维索引功能。.accessor<>接口旨在高效访问CPU张量上的数据。对于CUDA张量，等效的接口是packed_accessor64<>和packed_accessor32<>，它们生成具有64位或32位整数索引的Packed Accessors。

Packed Accessor与Accessor的根本区别在于，Packed Accessor将其大小和步长数据复制到其结构内部，而不是指向外部数据。这使得我们可以将其传递给CUDA核函数，并在其中使用其接口。

我们可以设计一个函数，使其接收Packed Accessors而不是指针。

__global__voidlltm_cuda_forward_kernel(
consttorch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>gates,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>old_cell,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>new_h,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>new_cell,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>input_gate,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>output_gate,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>candidate_cell)

让我们分解一下这里使用的模板。前两个参数 scalar_t 和 2 与常规的 Accessor 相同。参数 torch::RestrictPtrTraits 表示必须使用 __restrict__ 关键字。还要注意我们使用了 PackedAccessor32 变体，它将大小和步长存储在 int32_t 中。这一点很重要，因为使用 64 位变体 (PackedAccessor64) 可能会导致内核变慢。

函数声明变为：

template<typenamescalar_t>
__global__voidlltm_cuda_forward_kernel(
consttorch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>gates,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>old_cell,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>new_h,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>new_cell,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>input_gate,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>output_gate,
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>candidate_cell){
//batch index
constintn=blockIdx.y;
// column index
constintc=blockIdx.x*blockDim.x+threadIdx.x;
if(c<gates.size(2)){
input_gate[n][c]=sigmoid(gates[n][0][c]);
output_gate[n][c]=sigmoid(gates[n][1][c]);
candidate_cell[n][c]=elu(gates[n][2][c]);
new_cell[n][c]=
old_cell[n][c]+candidate_cell[n][c]*input_gate[n][c];
new_h[n][c]=tanh(new_cell[n][c])*output_gate[n][c];
}
}

实现变得更加易读！然后通过在主机函数中使用 .packed_accessor32<> 方法创建 Packed Accessors 来调用此函数。

std::vector<torch::Tensor>lltm_cuda_forward(
torch::Tensorinput,
torch::Tensorweights,
torch::Tensorbias,
torch::Tensorold_h,
torch::Tensorold_cell){
autoX=torch::cat({old_h,input},/*dim=*/1);
autogate_weights=torch::addmm(bias,X,weights.transpose(0,1));

constautobatch_size=old_cell.size(0);
constautostate_size=old_cell.size(1);

autogates=gate_weights.reshape({batch_size,3,state_size});
autonew_h=torch::zeros_like(old_cell);
autonew_cell=torch::zeros_like(old_cell);
autoinput_gate=torch::zeros_like(old_cell);
autooutput_gate=torch::zeros_like(old_cell);
autocandidate_cell=torch::zeros_like(old_cell);

constintthreads=1024;
constdim3blocks((state_size+threads-1)/threads,batch_size);

AT_DISPATCH_FLOATING_TYPES(gates.type(),"lltm_forward_cuda",([&]{
lltm_cuda_forward_kernel<scalar_t><<<blocks,threads>>>(
gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
old_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
new_h.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
new_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
input_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
output_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
candidate_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
}));

return{new_h,new_cell,input_gate,output_gate,candidate_cell,X,gates};
}

反向传播遵循类似的模式，我就不再赘述了：

template<typenamescalar_t>
__global__voidlltm_cuda_backward_kernel(
torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>d_old_cell,
torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>d_gates,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>grad_h,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>grad_cell,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>new_cell,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>input_gate,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>output_gate,
consttorch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits>candidate_cell,
consttorch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits>gate_weights){
//batch index
constintn=blockIdx.y;
// column index
constintc=blockIdx.x*blockDim.x+threadIdx.x;
if(c<d_gates.size(2)){
constautod_output_gate=tanh(new_cell[n][c])*grad_h[n][c];
constautod_tanh_new_cell=output_gate[n][c]*grad_h[n][c];
constautod_new_cell=
d_tanh(new_cell[n][c])*d_tanh_new_cell+grad_cell[n][c];


d_old_cell[n][c]=d_new_cell;
constautod_candidate_cell=input_gate[n][c]*d_new_cell;
constautod_input_gate=candidate_cell[n][c]*d_new_cell;

d_gates[n][0][c]=
d_input_gate*d_sigmoid(gate_weights[n][0][c]);
d_gates[n][1][c]=
d_output_gate*d_sigmoid(gate_weights[n][1][c]);
d_gates[n][2][c]=
d_candidate_cell*d_elu(gate_weights[n][2][c]);
}
}

std::vector<torch::Tensor>lltm_cuda_backward(
torch::Tensorgrad_h,
torch::Tensorgrad_cell,
torch::Tensornew_cell,
torch::Tensorinput_gate,
torch::Tensoroutput_gate,
torch::Tensorcandidate_cell,
torch::TensorX,
torch::Tensorgates,
torch::Tensorweights){
autod_old_cell=torch::zeros_like(new_cell);
autod_gates=torch::zeros_like(gates);

constautobatch_size=new_cell.size(0);
constautostate_size=new_cell.size(1);

constintthreads=1024;
constdim3blocks((state_size+threads-1)/threads,batch_size);

AT_DISPATCH_FLOATING_TYPES(X.type(),"lltm_backward_cuda",([&]{
lltm_cuda_backward_kernel<scalar_t><<<blocks,threads>>>(
d_old_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
d_gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
grad_h.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
grad_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
new_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
input_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
output_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
candidate_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
}));

autod_gate_weights=d_gates.reshape({batch_size,3*state_size});
autod_weights=d_gate_weights.t().mm(X);
autod_bias=d_gate_weights.sum(/*dim=*/0,/*keepdim=*/true);

autod_X=d_gate_weights.mm(weights);
autod_old_h=d_X.slice(/*dim=*/1,0,state_size);
autod_input=d_X.slice(/*dim=*/1,state_size);

return{d_old_h,d_input,d_weights,d_bias,d_old_cell,d_gates};
}

将 C++/CUDA 操作集成到 PyTorch 中

将我们支持 CUDA 的操作与 PyTorch 集成同样非常简单。如果你想编写一个 setup.py 脚本，它可能如下所示：

fromsetuptoolsimport setup
fromtorch.utils.cpp_extensionimport BuildExtension, CUDAExtension

setup(
    name='lltm',
    ext_modules=[
        CUDAExtension('lltm_cuda', [
            'lltm_cuda.cpp',
            'lltm_cuda_kernel.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

现在，我们使用 CUDAExtension() 而不是 CppExtension()。我们可以直接指定 .cu 文件以及 .cpp 文件——库会为您处理所有相关的繁琐操作。JIT 机制甚至更加简单：

fromtorch.utils.cpp_extensionimport load

lltm = load(name='lltm', sources=['lltm_cuda.cpp', 'lltm_cuda_kernel.cu'])

性能对比

我们的期望是通过使用 CUDA 并行化并融合代码中的逐元素操作，能够提升 LLTM 的性能。让我们看看这一期望是否成立。我们可以运行前面列出的代码来进行基准测试。之前最快的版本是基于 CUDA 的 C++ 代码：

Forward: 149.802 us | Backward 393.458 us

现在，使用我们自定义的CUDA内核：

Forward: 129.431 us | Backward 304.641 us

性能进一步提升！

结论

现在，您应该对 PyTorch 的 C++ 扩展机制及其使用动机有了一个全面的了解。您可以在这里找到本笔记中展示的代码示例。如果您有任何问题，请使用论坛。此外，如果您遇到任何问题，请务必查看我们的常见问题解答。关于为 AMD ROCm 编写扩展的博客可以在这里找到。

自定义 C++ 和 CUDA 扩展

动机与示例

编写 C++ 扩展

使用 setuptools 构建

编写 C++ 操作

前向传播

反向传播

绑定到 Python

使用您的扩展

性能对比

GPU 设备上的性能

JIT 编译扩展

编写混合 C++/CUDA 扩展

使用访问器

将 C++/CUDA 操作集成到 PyTorch 中

性能对比

结论

使用 `setuptools` 构建