Pytorch自定义C++/CUDA扩展

翻译自：官方文档

PyTorch 提供了大量与神经网络、张量代数、数据整理和其他操作。但是，我们有时会需要更加定制化的操作。例如，想要使用论文中找到的一种新型的激活函数，或者实现自己设计的算子。

在 PyTorch 中集成此类自定义操作的最简单方法是使用 Python 扩展这里概述的 Function 和 Module 。这里已经提供了自动微分的全部功能（无需编写计算微分的函数）以及 Python 的常用的表达是。但是，有时算子更适合用 C++ 实现。例如，某些算子可能需要非常快，因为它在模型中被非常频繁地调用，或者即使很少调用也非常耗时。另一个可能的原因是某些算子依赖于其他 C 或 C++ 库或。为了解决这种情况，PyTorch 提供了一种编写自定义 C++ 扩展的非常简单的方法。

C++ 扩展允许用户创建在 out-of-source 定义的 PyTorch 算子，即与 PyTorch 后端分离。这种方法不同于实现原生 PyTorch 操作的方式。 C++ 扩展旨在节省大量与将操作与 PyTorch 的后端集成相关的样板，同时为基于 PyTorch 的项目提供高度的灵活性。然而，一旦将算子定义为 C++ 扩展，将其转换为原生 PyTorch 函数主要是代码组织问题，如果决定将操作贡献给上游，则可以事后解决。

动机、实例与Python扩展实现

本文下面的部分将介绍一个编写和使用 C++（和 CUDA）扩展的实例。

假设我们提出了一种新的循环单元，这个循环单元类似于 LSTM，但不同之处在于它没有遗忘门，并使用指数线性单元 (ELU) 作为其内部激活函数。因为这个单元永远不会忘记，我们称之为 LLTM，或 Long-Long-Term-Memory 单元。

LLTM 与普通 LSTM 的两种不同之处非常重要，以至于我们无法为我们的目的配置 PyTorch 的 LSTM 算子，因此我们必须创建一个自定义算子。第一个也是最简单的方法——可能在所有情况下都是很好的第一步——是用 Python 在普通的 PyTorch 中实现我们想要的功能。为此，我们需要继承 torch.nn.Module 并实现 LLTM 的 forward。即：

import torch
import math
import torch.nn.functional as Fclass LLTM(torch.nn.Module):def __init__(self, input_features, state_size):super(LLTM, self).__init__()self.input_features = input_featuresself.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()def reset_parameters(self):stdv = 1.0 / math.sqrt(self.state_size)for weight in self.parameters():weight.data.uniform_(-stdv, +stdv)def forward(self, input, state):old_h, old_cell = stateX = 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_gatereturn new_h, new_cell

然后我们可以这样调用：

import torchbatch_size = 4
input_features = 12
state_size = 8X = 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 的融合版本。我们将从使用纯 C++ 编写它开始，使用为 PyTorch 的大部分后端提供支持的 ATen 库。然后，我们将通过将模型的一部分移动到 CUDA 内核以从 GPU 提供的大规模并行性中受益，从而进一步加快速度。

实现C++扩展

C++ 扩展有两种形式：它们可以使用

setuptools “ahead of time (AOT)” 构建，
或者通过 torch.utils.cpp_extension.load() “just in time (JIT)” 构建。

我们将逐个介绍。

aot编译扩展

对于“ahead of time”方式，我们通过编写一个 setup.py 脚本来构建我们的 C++ 扩展，该脚本使用 setuptools 来编译我们的 C++ 代码。

from setuptools import setup, Extension
from torch.utils import cpp_extensionsetup(name='lltm_cpp',ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['lltm.cpp'])],cmdclass={'build_ext': cpp_extension.BuildExtension})

在此代码中，CppExtension 是 setuptools.Extension 的一个方便的 wrapper，它传递正确的包含路径并将扩展的语言设置为 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::Tensor d_sigmoid(torch::Tensor z) {auto s = torch::sigmoid(z);return (1 - s) * s;
}

<torch/extension.h> 是 “一站式” 头文件，包含编写 C++ 扩展所需的所有 PyTorch 内容。这包括：

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

d_sigmoid() 的实现展示了如何使用 ATen API。 PyTorch 的张量和变量接口是从 ATen 库自动生成的，因此我们可以或多或少地将 Python 实现 1:1 转换为 C++。我们所有计算的主要数据类型是 torch::Tensor。可以在这里查看其完整的 API。另请注意，我们可以包含 <iostream> 等 C 或 C++ 头文件，并支持 C++11 的全部功能。

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

示例如下，使用：

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

而不要：

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

这时 nvcc 的一个bug，目前仍是 open 的 issue，完整解决的代码示例在这里。

前向传播

下面，我们给出完整前向传播的 C++ 实现：

#include <vector>std::vector<at::Tensor> lltm_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell) {auto X = torch::cat({old_h, input}, /*dim=*/1);auto gate_weights = torch::addmm(bias, X, weights.transpose(0, 1));auto gates = gate_weights.chunk(3, /*dim=*/1);auto input_gate = torch::sigmoid(gates[0]);auto output_gate = torch::sigmoid(gates[1]);auto candidate_cell = torch::elu(gates[2], /*alpha=*/1.0);auto new_cell = old_cell + candidate_cell * input_gate;auto new_h = torch::tanh(new_cell) * output_gate;return {new_h,new_cell,input_gate,output_gate,candidate_cell,X,gate_weights};
}

反向传播

C++ 扩展 API 目前没有为我们提供自动生成反向传播函数的方法（之前提到 Python 可以）。因此，我们还必须自己实现 LLTM 的反向传播，它计算损失关于前向传播的每个输入的导数。最终，我们将前向和后向函数都放入 torch.autograd.Function 中，来创建一个的 Python binding。反向函数稍微复杂一些，因此我们不会深入研究代码（如果有兴趣，可以阅读 Alex Graves 的论文以获取更多信息）：

// tanh'(z) = 1 - tanh^2(z)
torch::Tensor d_tanh(torch::Tensor z) {return 1 - z.tanh().pow(2);
}// elu'(z) = relu'(z) + { alpha * exp(z) if (alpha * (exp(z) - 1)) < 0, else 0}
torch::Tensor d_elu(torch::Tensor z, torch::Scalar alpha = 1.0) {auto e = z.exp();auto mask = (alpha * (e - 1)) < 0;return (z > 0).type_as(z) + mask.type_as(z) * (alpha * e);
}std::vector<torch::Tensor> lltm_backward(torch::Tensor grad_h,torch::Tensor grad_cell,torch::Tensor new_cell,torch::Tensor input_gate,torch::Tensor output_gate,torch::Tensor candidate_cell,torch::Tensor X,torch::Tensor gate_weights,torch::Tensor weights) {auto d_output_gate = torch::tanh(new_cell) * grad_h;auto d_tanh_new_cell = output_gate * grad_h;auto d_new_cell = d_tanh(new_cell) * d_tanh_new_cell + grad_cell;auto d_old_cell = d_new_cell;auto d_candidate_cell = input_gate * d_new_cell;auto d_input_gate = candidate_cell * d_new_cell;auto gates = 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]);auto d_gates =torch::cat({d_input_gate, d_output_gate, d_candidate_cell}, /*dim=*/1);auto d_weights = d_gates.t().mm(X);auto d_bias = d_gates.sum(/*dim=*/0, /*keepdim=*/true);auto d_X = d_gates.mm(weights);const auto state_size = grad_h.size(1);auto d_old_h = d_X.slice(/*dim=*/1, 0, state_size);auto d_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.cppsetup.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.pthInstalled ~/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，然后使用相同的编译器构建扩展。

构建扩展后，我们可以使用在 setup.py 脚本中指定的名称将其简单地导入 Python。请务必先导入 torch，因为这将解析动态链接器必须看到的一些符号：

In [1]: import torch
In [2]: import lltm_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 instanceforward(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 的一等公民：

import math
import torch# Our module!
import lltm_cppclass LLTMFunction(torch.autograd.Function):@staticmethoddef forward(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@staticmethoddef backward(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 = outputsreturn d_input, d_weights, d_bias, d_old_h, d_old_cellclass LLTM(torch.nn.Module):def __init__(self, input_features, state_size):super(LLTM, self).__init__()self.input_features = input_featuresself.state_size = state_sizeself.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()def reset_parameters(self):stdv = 1.0 / math.sqrt(self.state_size)for weight in self.parameters():weight.data.uniform_(-stdv, +stdv)def forward(self, input, state):return LLTMFunction.apply(input, self.weights, self.bias, *state)

性能对比

现在我们可以从 PyTorch 使用和调用我们的 C++ 代码，我们可以运行一个小型基准测试，看看我们通过用 C++ 重写我们的操作获得了多少性能。我们将前向和反向运行 LLTM 几次并测量持续时间：

import timeimport torchbatch_size = 16
input_features = 32
state_size = 128X = 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() - startstart = time.time()(new_h.sum() + new_C.sum()).backward()backward += time.time() - startprint('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

GPU上的性能

关于 PyTorch 的 ATen 后端的一个奇妙的事实是它抽象了我们正在运行的计算设备。这意味着我们为 CPU 编写的相同代码也可以在 GPU 上运行，并且各个操作将相应地分派给 GPU 优化的实现。对于矩阵乘法（如 mm 或 addmm）等某些操作，这是一个巨大的胜利。让我们看看使用 CUDA 张量运行 C++ 代码可以获得多少性能。不需要对我们的实现进行任何更改，我们只需要将我们的张量从 Python 放入 GPU 内存中，在创建时添加 device=cuda_device 参数或在创建后使用 .to(cuda_device) ：

import torchassert torch.cuda.is_available()
cuda_device = torch.device("cuda")  # device object representing GPUbatch_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() - startstart = time.time()(new_h.sum() + new_C.sum()).backward()torch.cuda.synchronize()backward += time.time() - startprint('Forward: {:.3f} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))

再次将我们的普通 PyTorch 代码与现在都在 CUDA 设备上运行的 C++ 版本进行比较，我们再次看到了性能提升。对于 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 或 just in time (JIT)。讲完了前者，我们再来谈谈后者。 JIT 编译机制通过调用 PyTorch API 中名为 torch.utils.cpp_extension.load() 的函数，为我们提供了一种即时编译和加载扩展的方法。对于 LLTM：

from torch.utils.cpp_extension import loadlltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

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

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

如果想看到这个详细的过程，可以给 cpp_extension.load() 传递参数 verbose=True ：

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，这具有会特别有效，因为顺序中有大量的逐点操作，它们都可以在单个 CUDA 核中融合和并行化。让我们看看我们如何编写这样的 CUDA 核并使用这种扩展机制将其与 PyTorch 集成。

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

我们将从 C++ 文件开始，我们将其称为 lltm_cuda.cpp，如下：

#include <torch/extension.h>#include <vector>// CUDA forward declarationsstd::vector<torch::Tensor> lltm_cuda_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell);std::vector<torch::Tensor> lltm_cuda_backward(torch::Tensor grad_h,torch::Tensor grad_cell,torch::Tensor new_cell,torch::Tensor input_gate,torch::Tensor output_gate,torch::Tensor candidate_cell,torch::Tensor X,torch::Tensor gate_weights,torch::Tensor weights);// 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::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell) {CHECK_INPUT(input);CHECK_INPUT(weights);CHECK_INPUT(bias);CHECK_INPUT(old_h);CHECK_INPUT(old_cell);return lltm_cuda_forward(input, weights, bias, old_h, old_cell);
}std::vector<torch::Tensor> lltm_backward(torch::Tensor grad_h,torch::Tensor grad_cell,torch::Tensor new_cell,torch::Tensor input_gate,torch::Tensor output_gate,torch::Tensor candidate_cell,torch::Tensor X,torch::Tensor gate_weights,torch::Tensor weights) {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);return lltm_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 <typename scalar_t>
__device__ __forceinline__ scalar_t sigmoid(scalar_t z) {return 1.0 / (1.0 + exp(-z));
}

在这里，我们看到了我刚刚描述的头文件，以及我们使用 CUDA 特定语法（如 __device__ 和 __forceinline__ ）和函数（如 exp）。我们继续使用一些辅助函数：

template <typename scalar_t>
__device__ __forceinline__ scalar_t d_sigmoid(scalar_t z) {const auto s = sigmoid(z);return (1.0 - s) * s;
}template <typename scalar_t>
__device__ __forceinline__ scalar_t d_tanh(scalar_t z) {const auto t = tanh(z);return 1 - (t * t);
}template <typename scalar_t>
__device__ __forceinline__ scalar_t elu(scalar_t z, scalar_t alpha = 1.0) {return fmax(0.0, z) + fmin(0.0, alpha * (exp(z) - 1.0));
}template <typename scalar_t>
__device__ __forceinline__ scalar_t d_elu(scalar_t z, scalar_t alpha = 1.0) {const auto e = exp(z);const auto d_relu = z < 0.0 ? 0.0 : 1.0;return d_relu + (((alpha * (e - 1.0)) < 0.0) ? (alpha * e) : 0.0);
}

现在要实际实现一个函数，我们还需要两件事：一个执行我们不希望显式手动编写的操作并调用 CUDA 内核的函数，然后是我们想要加速的部分的实际 CUDA 内核 . 对于前向传递，第一个函数应该如下所示：

std::vector<torch::Tensor> lltm_cuda_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell) {auto X = torch::cat({old_h, input}, /*dim=*/1);auto gates = torch::addmm(bias, X, weights.transpose(0, 1));const auto batch_size = old_cell.size(0);const auto state_size = old_cell.size(1);auto new_h = torch::zeros_like(old_cell);auto new_cell = torch::zeros_like(old_cell);auto input_gate = torch::zeros_like(old_cell);auto output_gate = torch::zeros_like(old_cell);auto candidate_cell = torch::zeros_like(old_cell);const int threads = 1024;const dim3 blocks((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()) {case torch::ScalarType::Double:return function<double>(tensor.data<double>());case torch::ScalarType::Float:return function<float>(tensor.data<float>());...
}

AT_DISPATCH_FLOATING_TYPES 的目的是为我们处理这个调度。它需要一个类型（在我们的例子中是 gates.type()）、一个名称（用于显示错误消息）和一个 lambda 函数。在这个 lambda 函数中，类型别名 scalar_t 可用，并被定义为张量在运行时在该 context 中实际存在的类型。因此，如果我们有一个模板函数（就是我们的 CUDA 核函数），我们可以用这个 scalar_t 别名实例化它，然后调用正确的函数。在这种情况下，我们还希望检索张量的数据指针作为该 scalar_t 类型的指针。如果想分派所有类型而不仅仅是浮点类型（Float 和 Double），可以使用 AT_DISPATCH_ALL_TYPES。

请注意，我们使用纯 ATen 执行一些操作。这些操作仍将在 GPU 上运行，但使用 ATen 的默认实现。这是有道理的，因为 ATen 将使用高度优化的例程来处理矩阵乘法（例如 addmm）或卷积，这部分很难实现和改进。

至于内核启动本身，我们在此指定每个 CUDA block 有 1024 个线程，并且整个 GPU 网格被分成多个 1 x 1024 线程块，以每个组件一个线程填充我们的矩阵。例如，如果我们的 state_size 大小为 2048，batch_size 大小为 4，我们将启动总共 4 x 2 = 8 个块，每个 1024 个线程。如果不了解 CUDA“block”或“grid”，那么可以参考阅读有关 CUDA 的介绍性读物。

实际的 CUDA 核相当简单（如果有 GPU 编程的经验）：

template <typename scalar_t>
__global__ void lltm_cuda_forward_kernel(const scalar_t* __restrict__ gates,const scalar_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_t state_size) {const int column = blockIdx.x * blockDim.x + threadIdx.x;const int index = blockIdx.y * state_size + column;const int gates_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 循环超过一百万个串行元素来做到这一点，就会明白为什么这会快得多。

使用accessors

可以在 CUDA 核中看到我们直接处理对应类型的指针。事实上，直接在 cuda 核中使用高级类型不可知的张量是非常低效的。然而，这是以易用性和可读性为代价的，尤其是对于高维数据。在我们的示例中，我们知道例如连续门张量具有 3 个维度：

batch，大小为 batch_size，步长为 3*state_size
row，大小为 3，步长为 state_size
index，大小为 state_size ，步长为 1

那么我们如何访问内核中的元素 gates[n][row][column] 呢？事实证明，可以通过一些简单的算术来访问对应位置的元素。

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

这个表达式需要明确知道步长，并在其参数中传递给核函数。可以看到，在核函数接受多个不同大小的张量的情况下，最终会得到一个很长的参数列表。对我们来说幸运的是，ATen 提供了通过单一动态检查创建的访问器，张量是维度的类型和数量。然后访问器公开一个 API 以有效地访问张量元素，而无需转换为单个指针：

torch::Tensor foo = torch::rand({12, 12});// assert foo is 2-dimensional and holds floats.
auto foo_a = foo.accessor<float,2>();
float trace = 0;for(int i = 0; i < foo_a.size(0); i++) {// use the accessor foo_a to get tensor data.trace += foo_a[i][i];
}

Accessor 对象有一个相对高层的接口，具有 .size() 和 .stride() 方法以及多维索引。 .accessor<> 接口旨在有效地访问 cpu 张量上的数据。对应于 cuda 张量的是 packed_accessor64<> 和 packed_accessor32<>，它们生成具有 64 位或 32 位整数索引的打包 accessor。

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

我们可以设计一个使用 Packed Accessors 而不是指针的函数。

__global__ void lltm_cuda_forward_kernel(const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> gates,const torch::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 <typename scalar_t>
__global__ void lltm_cuda_forward_kernel(const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> gates,const torch::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 indexconst int n = blockIdx.y;// column indexconst int c = 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];}
}

该实现更具可读性。然后通过在 host 函数中使用 .packed_accessor32<> 方法创建打包访问器来调用此函数。

std::vector<torch::Tensor> lltm_cuda_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell) {auto X = torch::cat({old_h, input}, /*dim=*/1);auto gate_weights = torch::addmm(bias, X, weights.transpose(0, 1));const auto batch_size = old_cell.size(0);const auto state_size = old_cell.size(1);auto gates = gate_weights.reshape({batch_size, 3, state_size});auto new_h = torch::zeros_like(old_cell);auto new_cell = torch::zeros_like(old_cell);auto input_gate = torch::zeros_like(old_cell);auto output_gate = torch::zeros_like(old_cell);auto candidate_cell = torch::zeros_like(old_cell);const int threads = 1024;const dim3 blocks((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 <typename scalar_t>
__global__ void lltm_cuda_backward_kernel(torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> d_old_cell,torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> d_gates,const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> grad_h,const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> grad_cell,const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> new_cell,const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> input_gate,const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> output_gate,const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> candidate_cell,const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> gate_weights) {//batch indexconst int n = blockIdx.y;// column indexconst int c = blockIdx.x * blockDim.x + threadIdx.x;if (c < d_gates.size(2)){const auto d_output_gate = tanh(new_cell[n][c]) * grad_h[n][c];const auto d_tanh_new_cell = output_gate[n][c] * grad_h[n][c];const auto d_new_cell =d_tanh(new_cell[n][c]) * d_tanh_new_cell + grad_cell[n][c];d_old_cell[n][c] = d_new_cell;const auto d_candidate_cell = input_gate[n][c] * d_new_cell;const auto d_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::Tensor grad_h,torch::Tensor grad_cell,torch::Tensor new_cell,torch::Tensor input_gate,torch::Tensor output_gate,torch::Tensor candidate_cell,torch::Tensor X,torch::Tensor gates,torch::Tensor weights) {auto d_old_cell = torch::zeros_like(new_cell);auto d_gates = torch::zeros_like(gates);const auto batch_size = new_cell.size(0);const auto state_size = new_cell.size(1);const int threads = 1024;const dim3 blocks((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>());}));auto d_gate_weights = d_gates.reshape({batch_size, 3*state_size});auto d_weights = d_gate_weights.t().mm(X);auto d_bias = d_gate_weights.sum(/*dim=*/0, /*keepdim=*/true);auto d_X = d_gate_weights.mm(weights);auto d_old_h = d_X.slice(/*dim=*/1, 0, state_size);auto d_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 脚本，如下所示：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtensionsetup(name='lltm',ext_modules=[CUDAExtension('lltm_cuda', ['lltm_cuda.cpp','lltm_cuda_kernel.cu',])],cmdclass={'build_ext': BuildExtension})

我们现在使用 CUDAExtension() 代替 CppExtension()。我们可以只指定 .cu 文件和 .cpp 文件——该库会为我们解决所有这些麻烦。 JIT 机制则更简单：

from torch.utils.cpp_extension import loadlltm = 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++ 扩展机制以及使用它们的动机有了很好的了解。可以在这里找到本文中展示的代码示例。如果有任何问题，请使用论坛，或者看常见问题的 FAQ。