自定义 C++ 和 CUDA 扩展¶
Created On: Apr 26, 2018 | Last Updated: Mar 19, 2025 | Last Verified: Nov 05, 2024
作者: Peter Goldsborough <https://www.goldsborough.me/>
警告
本教程从 PyTorch 2.4 起已被弃用。请参阅 PyTorch自定义操作符 了解最新的关于扩展 PyTorch 的自定义 C++/CUDA 扩展指南。
PyTorch 提供了大量涉及神经网络、任意张量代数、数据处理等方面的操作。然而，您可能仍然需要更为定制化的操作。例如，您可能希望使用在论文中发现的一种新型激活函数，或者实现作为研究一部分开发的某个操作。
在 PyTorch 中集成这样一个定制化操作的最简单方法是通过在 Python 中扩展 Function 和 Module 来实现，如 这里 <https://pytorch.org/docs/master/notes/extending.html> 所述。这种方式赋予了您全自动求导的强大功能（无需编写导数函数）以及 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 的前向传播。这看起来大致如下：
class LLTM(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()
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 = 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
然后我们可以像预期那样使用它：
import torch
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 <https://developer.nvidia.com/cudnn>、Intel MKL <https://software.intel.com/en-us/mkl> 或 NNPACK <https://github.com/Maratyszcza/NNPACK> 这样的库，像上面这样的 PyTorch 代码通常已经足够快。然而，我们也可以看到，在某些情况下仍然有进一步优化的空间。最显而易见的原因是 PyTorch 不了解您正在实现的*算法*。它只知道您用于构成算法的个别操作。因此，PyTorch 必须逐个执行您的操作。由于每次调用操作的实现（或*内核*）可能涉及 CUDA 内核的启动，因此这种启动开销可能在多次函数调用中显得较为重要。此外，运行代码的 Python 解释器本身也可能会减慢程序。
因此，加速代码的一个明确的方法是将部分逻辑重写为 C++（或 CUDA），并*融合*某些特定的操作组。融合意味着将许多功能的实现组合为一个功能，这样可以减少内核启动并在数据流的全局视图方面进行其他优化。
让我们看看如何使用 C++ 扩展来实现 LLTM 的*融合*版本。我们将首先用简单的 C++ 编写它，通过 ATen <https://github.com/zdevito/ATen> 库（PyTorch 后端的核心部分）来实现，并看看它是如何轻松地将我们的 Python 代码翻译为实际应用的。然后我们将进一步提速，通过将模型的部分迁移到 CUDA 内核，以利用 GPU 提供的大规模并行计算能力。
编写 C++ 扩展¶
C++ 扩展有两种形式：可以使用 setuptools 进行“预构建”，或者使用 torch.utils.cpp_extension.load() 进行“即时加载”。我们将先讨论第一种方法，然后再说明另一种。
使用 setuptools 进行构建¶
对于“预构建”方式，我们通过编写一个使用 setuptools 编译 C++ 代码的“setup.py”脚本来构建 C++ 扩展。对于 LLTM，它看起来非常简单，如下所示：
from setuptools import setup, Extension
from torch.utils import 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::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 上使用 CUDA-11.5 的 nvcc 解析 torch/extension.h 时会出现内部编译器错误。为了规避此问题，请将 Python 绑定逻辑移动到纯 C++ 文件中。示例用法：
#include <ATen/ATen.h>
at::Tensor SigmoidAlphaBlendForwardCuda(....)
代替：
#include <torch/extension.h>
torch::Tensor SigmoidAlphaBlendForwardCuda(...)
Currently open issue for nvcc bug here. Complete workaround code example here.
前向传播¶
接下来，我们可以将整个前向传播移植到 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 并未提供一种方法来自动为我们生成反向函数。因此，我们还必须实现 LLTM 的反向传播，该过程计算损失对前向传播每个输入的导数。最终，我们会将前向和反向函数封装到 torch.autograd.Function 中，以创建一个优雅的 Python 绑定。反向函数稍微复杂一些，因此我们不会深入探讨代码（如果您感兴趣，可以阅读 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.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]: 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 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 的一级公民：
import math
import torch
# Our module!
import lltm_cpp
class LLTMFunction(torch.autograd.Function):
@staticmethod
def 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
@staticmethod
def 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 = outputs
return d_input, d_weights, d_bias, d_old_h, d_old_cell
class LLTM(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()
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 time
import torch
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++ 代码可以获得多少性能提升。无需对实现进行更改，我们只需要通过在创建时添加 device=cuda_device 参数，或者在创建后使用 .to(cuda_device)，将张量放入 GPU 内存：
import torch
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，这将变得如此简单：
from torch.utils.cpp_extension import 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 来说，这特别有效，因为有大量连续的逐点操作，可以在单个 CUDA 核心中融合并并行化。让我们看看如何编写这样的 CUDA 核心，并使用此扩展机制将其与 PyTorch 集成。
编写 CUDA 扩展的一般策略是首先编写一个定义将从 Python 调用的函数的 C++ 文件，并使用 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 declarations
std::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``）。注意：:mod:`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``是可用的，并定义为在相应上下文中张量在运行时实际的类型。因此，如果我们有一个模板函数（我们的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的介绍文章 <https://devblogs.nvidia.com/even-easier-introduction-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``循环序列处理一百万个元素，就可以理解为什么这会快得多。
使用访问器¶
您可以在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]
此外，这种表达式不仅冗长，还需要显式知道步幅，因此必须通过参数传递给内核函数。您可以看到，对于接受具有不同尺寸的多个张量的内核函数，您最终会得到一个非常长的参数列表。
幸运的是，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];
}
访问器对象具有相对高级的接口，带有``.size()``和``.stride()``方法以及多维索引。用于访问CPU张量的``.accessor<>``接口的等效项是``packed_accessor64<>``和``packed_accessor32<>``，它们分别生成具有64位或32位整数索引的打包访问器。
与普通访问器的根本区别在于，打包访问器将大小或步幅的数据复制到其结构内，而不是指向它。这使我们能够将其传递给CUDA内核函数并在其中使用其接口。
我们可以设计一个接受打包访问器而不是指针的函数。
__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``与普通访问器相同。参数``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 index
const int n = blockIdx.y;
// column index
const 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];
}
}
实现变得更易读！然后通过在主机函数中使用``.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 index
const int n = blockIdx.y;
// column index
const 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, CUDAExtension
setup(
name='lltm',
ext_modules=[
CUDAExtension('lltm_cuda', [
'lltm_cuda.cpp',
'lltm_cuda_kernel.cu',
])
],
cmdclass={
'build_ext': BuildExtension
})
我们现在使用:func:CUDAExtension`而不是:func:`CppExtension。我们可以简单地指定``.cu``文件以及``.cpp``文件——库会为您处理所有相关的复杂工作。JIT机制甚至更简单：
from torch.utils.cpp_extension import 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++扩展机制有了一个很好的概述并了解使用它们的动机。您可以在`这里 <https://github.com/pytorch/extension-cpp>`_找到本笔记中展示的代码示例。如果您有问题，请使用`论坛 <https://discuss.pytorch.org>`_。如果您遇到任何问题，务必查看我们的`FAQ <https://pytorch.org/cppdocs/notes/faq.html>`_。关于为AMD ROCm编写扩展的博客请参阅`这里 <https://rocm.blogs.amd.com/artificial-intelligence/cpp-extn/readme.html>`_。
