自定义C++和CUDA扩展¶
创建于：2018年4月26日 | 最后更新于：2024年7月22日 | 最后验证于：2024年11月5日
作者: 彼得·戈兹伯勒
警告
本教程在 PyTorch 2.4 版本中已弃用。请参阅 PyTorch 自定义操作符 以获取有关使用自定义 C++/CUDA 扩展扩展 PyTorch 的最新指南。
PyTorch 提供了大量的与神经网络相关的操作、任意张量代数运算、数据处理以及其他用途的操作。然而，您可能仍然需要一个更加定制化的操作。例如，您可能想要使用一篇论文中发现的新激活函数，或者实现自己研究中开发的操作。
将这种自定义操作集成到PyTorch中最简单的方法是通过扩展Function和Module用Python编写，如这里所述。这使你能够充分利用自动微分的优势（无需编写导数函数）以及Python通常的表达能力。然而，有时你的操作可能更适合用C++实现。例如，你的代码可能需要非常快，因为它在模型中被频繁调用，或者即使很少调用也非常耗时。另一个合理的原因是它依赖于或与其他C或C++库交互。为了解决这些问题，PyTorch提供了一种非常简便的方法来编写自定义的C++扩展。
C++ 扩展是一种机制，我们开发它以允许用户（您）创建定义在 源外 的 PyTorch 操作，即与 PyTorch 后端分离的操作。这种方法与原生 PyTorch 操作的实现方式不同。C++ 扩展旨在让您免去许多将操作与 PyTorch 后端集成的样板代码，同时为您提供高度的灵活性来支持您的基于 PyTorch 的项目。 不过，一旦您将操作定义为 C++ 扩展，将其转换为原生 PyTorch 函数主要是一个代码组织问题，如果您决定上游贡献您的操作，您可以事后解决这个问题。
动机和示例¶
本说明的其余部分将通过一个实际示例来介绍如何编写和使用C++（以及CUDA）扩展。如果你正身处险境或者有人会在今天结束前因为你没有完成那个操作而解雇你，你可以跳过这一部分，直接前往下一节的实现细节。
让我们说你提出了一种新的循环单元，发现它的性能优于当前最先进的技术。这种循环单元类似于LSTM，但不同之处在于它缺少一个忘记门，并且使用指数线性单元（ELU）作为其内部激活函数。由于这个单元永远不会忘记，我们将称它为LLTM，或长长久久记忆单元。
LLTMs与vanilla LSTMs的不同之处足够显著，以至于我们不能配置PyTorch的LSTMCell来满足我们的需求，所以我们需要创建一个自定义单元。对于这一点——并且在所有情况下，这很可能是一个好的第一步——我们可以使用纯Python在PyTorch中实现我们所需的功能。为此，我们需要继承
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、Intel MKL 或 NNPACK 等库提供支持，因此像上面那样的PyTorch代码通常已经足够快了。然而，我们也可以看到，在某些情况下，仍有进一步提高性能的空间。最明显的原因是，PyTorch并不了解你正在实现的算法。它只知道你用来组成算法的各个操作。因此，PyTorch必须逐个执行你的操作，一个接一个地进行。由于每个操作的实现（或内核）调用，可能涉及启动CUDA内核，都会有一定的开销，因此在许多函数调用中，这种开销可能会变得显著。此外，运行我们代码的Python解释器本身也可能减慢我们的程序。
因此，加快速度的一个确定方法是用C++（或CUDA）重写部分代码，并将特定组的操作进行 融合。融合意味着将许多函数的实现合并到一个函数中，从而受益于较少的内核启动以及其他我们可以通过对全局数据流动增加可见性来进行的优化。
让我们看看我们如何使用C++扩展来实现LLTM的融合版本。我们将从使用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++代码绑定的方式，
管理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(...)
当前打开的nvcc错误问题 这里。 完整的工作示例代码 这里。
前向传播¶
接下来，我们可以将整个前向传播过程移植到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 绑定。反向函数稍微复杂一些，所以我们不会深入探讨代码（如果你感兴趣，亚历克斯·格雷夫斯的论文 是了解更多信息的好读物）：
// 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/
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
A small note on 编译器：由于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++代码能获得多少性能提升。我们的实现不需要任何更改，我们只需要在Python中将张量放入GPU内存中，可以通过在创建时添加device=cuda_device参数或者创建后使用.to(cuda_device)来实现:
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))
再次比较我们的plain PyTorch代码与C++版本，现在两者都在CUDA设备上运行，我们又一次看到了性能提升。对于Python/PyTorch：
Forward: 187.719 us | Backward 410.815 us
And C++/ATen:
Forward: 149.802 us | Backward 393.458 us
与非CUDA代码相比，这是一个整体上的巨大加速。然而，我们可以通过编写自定义CUDA内核来进一步提升我们的C++代码性能，这我们很快就会探讨。在此之前，让我们讨论另一种构建C++扩展的方法。
JIT 编译扩展¶
之前，我提到过有两种方法可以构建C++扩展：使用setuptools或即时编译（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集成。
The general strategy for writing a 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::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>());
...
}
The purpose of AT_DISPATCH_FLOATING_TYPES is to take care of this dispatch for us. It takes a type (gates.type() in our case), a name (for error messages) and a lambda function. Inside this lambda function, the type alias scalar_t is available and is defined as the type that the tensor actually is at runtime in that context. As such, if we have a template function (which our CUDA kernel will be), we can instantiate it with this scalar_t alias, and the correct function will be called. In this case, we also want to retrieve the data pointers of the tensors as pointers of that scalar_t type. If you wanted to dispatch over all types and not just floating point types (Float and Double), you can use 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 <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];
}
_accessor对象具有相对较高的接口级别，包括.size()和
.stride()方法以及多维索引。_cpu张量的_接口_设计用于高效访问数据。cuda张量的等效接口是packed_accessor64<>和packed_accessor32<>，它们生成带有64位或32位整数索引的_packed _访问器。
与 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与常规访问器相同。第三个参数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
})
Instead of CppExtension(), 我们现在使用 CUDAExtension()。我们只需指定 .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
更多性能提升！