pytorch 的C++扩展
pytorch: cpp_extension
项目地址：https://github.com/firstelfin/torch_extension
深度学习中，我们常使用pytorch的python接口实现一些非官方实现的逻辑，我们可能会拼接一些官方的算子组成一个自己的新算子，这样做简单高效，但是在运行中计算性能可能会较差(多个算子叠加涉及到： PyTorch 不知道您正在实现的算法。它只知道您用来组成算法的各个操作。因此，PyTorch 必须一个接一个地单独执行您的操作。由于每个单独调用实现（或kernel ) 可能涉及启动 CUDA 内核的操作，具有一定的开销，这种开销可能在许多函数调用中变得显著。此外，运行我们代码的 Python 解释器本身可能会减慢我们的程序。)；在某些创新中，我们甚至无法基于常规官方算子进行开发，必须要自己书写整个逻辑。这时，如果使用python实现，我们可能会纠结其计算速度较低，形成计算瓶颈，所以这时候使用C++的扩展就非常重要了。
一、书写C++ Extension
C++ 扩展有两种形式：它们可以使用“预编译”构建 setuptools，或者通过 jit 即时编译 torch.utils.cpp_extension.load(). 我们将从第一种方法开始，稍后讨论后者。
这里我们希望实现一种新的LSTM(没有遗忘门)，称之为LLTM。
书写setup.py文件
对于“提前”风格，我们通过编写一个 setup.py脚本来构建我们的 C++ 扩展，该脚本使用 setuptools 来编译我们的 C++ 代码。对于 LLTM，它看起来很简单：
from setuptools import setup, Extension
from torch.utils import cpp_extension
setup(name='lltm_cpp',
ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['csrc/lltm.cpp'])],
cmdclass={'build_ext': cpp_extension.BuildExtension})
在此代码中，CppExtension是一个 setuptools.Extension的封装，它传递正确的路径并将扩展的语言设置为 C++。等效的setuptools 代码将是：
Extension(
name='lltm_cpp',
sources=['csrc/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>是包含所有必要 PyTorch bits以编写 C++ 扩展的一站式头文件。这包括：
- ATen 库，这是我们用于张量计算的主要 API，
- pybind11，这是我们为 C++ 代码创建 Python 绑定的方式，
- 管理 ATen 和 pybind11 之间交互细节的标头。
注：这里展示了如何使用ATen接口实现d_sigmoid. PyTorch 的张量和变量接口是从 ATen 库自动生成的，因此我们可以或多或少地将我们的 Python 实现 1:1 转换为 C++。我们所有计算的主要数据类型将是
torch::Tensor. 它的完整 API 可以在这里查看。另请注意，我们可以包含任何其他 C 或 C++ 头文件如<iostream>——我们拥有 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(...)
前向过程
#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)); // 矩阵乘法 X * weights + bias
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 的后向传递，它计算损失相对于前向传递的每个输入的导数。最终，我们会将 forward 和 backward 函数放入 atorch.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。pytorch扩展构建将其定义为您在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 /workspace/torch_extension-master/build
creating /workspace/torch_extension-master/build/temp.linux-x86_64-3.7
creating /workspace/torch_extension-master/build/temp.linux-x86_64-3.7/csrc
Emitting ninja build file /workspace/torch_extension-master/build/temp.linux-x86_64-3.7/build.ninja...
Compiling objects...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/1] c++ -MMD -MF /workspace/torch_extension-master/build/temp.linux-x86_64-3.7/csrc/lltm.o.d -pthread -B /opt/conda/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/opt/conda/lib/python3.7/site-packages/torch/include -I/opt/conda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I/opt/conda/lib/python3.7/site-packages/torch/include/TH -I/opt/conda/lib/python3.7/site-packages/torch/include/THC -I/opt/conda/include/python3.7m -c -c /workspace/torch_extension-master/csrc/lltm.cpp -o /workspace/torch_extension-master/build/temp.linux-x86_64-3.7/csrc/lltm.o -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=lltm_cpp -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14
cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++
In file included from /opt/conda/lib/python3.7/site-packages/torch/include/ATen/Parallel.h:140:0,
from /opt/conda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include/torch/utils.h:3,
from /opt/conda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include/torch/nn/cloneable.h:5,
from /opt/conda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include/torch/nn.h:3,
from /opt/conda/lib/python3.7/site-packages/torch/include/torch/csrc/api/include/torch/all.h:13,
from /opt/conda/lib/python3.7/site-packages/torch/include/torch/extension.h:4,
from /workspace/torch_extension-master/csrc/lltm.cpp:1:
/opt/conda/lib/python3.7/site-packages/torch/include/ATen/ParallelOpenMP.h:83:0: warning: ignoring #pragma omp parallel [-Wunknown-pragmas]
#pragma omp parallel for if ((end - begin) >= grain_size)
creating build/lib.linux-x86_64-3.7
g++ -pthread -shared -B /opt/conda/compiler_compat -L/opt/conda/lib -Wl,-rpath=/opt/conda/lib -Wl,--no-as-needed -Wl,--sysroot=/ /workspace/torch_extension-master/build/temp.linux-x86_64-3.7/csrc/lltm.o -L/opt/conda/lib/python3.7/site-packages/torch/lib -lc10 -ltorch -ltorch_cpu -ltorch_python -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
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
creating /opt/conda/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 /opt/conda/lib/python3.7/site-packages
Adding lltm-cpp 0.0.0 to easy-install.pth file
Installed /opt/conda/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，然后使用相同的编译器构建扩展。
使用这种扩展算子前我们必须先导入import torch， 再import lltm_cpp.
In [1]: import torch
In [2]: import lltm_cpp
In [3]: lltm_cpp.forward
Out[3]: <built-in method forward of PyCapsule object at ...>
如果使用help方法查看这个算子，它会匹配到C++的代码：
In[4] help(lltm_cpp.forward) # 慎用，你可能会卡死
Help on built-in function forward in module lltm_cpp:
forward(...) method of builtins.PyCapsule instance
forward(arg0: at::Tensor, arg1: at::Tensor, arg2: at::Tensor, arg3: at::Tensor, arg4: at::Tensor) -> List[at::Tensor]
LLTM forward
由于我们现在可以从 Python 调用我们的 C++ 函数，我们可以用 PyTorch 的torch.autograd.Function包装它们并用torch.nn.Module使它们成为基本算子：
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++ 实现的，因此预计会很快。尽管如此，这是一个好的开始。
本人docker容器测试：
rnn1 Forward: 22.805 s | Backward 30.756 s（PY）
rnn2 Forward: 18.791 s | Backward 51.717 s (C)
GPU上的性能
关于 PyTorch 的ATen后端的一个奇妙的事实是它抽象了您正在运行的计算设备。这意味着我们为 CPU 编写的相同代码也可以在 GPU 上运行，并且各个操作将相应地分派给 GPU 优化的实现。对于某些操作，如矩阵乘法（如mmor 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))
再次将我们的普通 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++ 扩展的方法。
自己代码的测试结果为：
rnn1 Forward: 18.546 s | Backward 27.655 s (PY)
rnn2 Forward: 14.757 s | Backward 49.820 s (C)为什么这里C++代码的反向传播慢很多呢？
注：以上代码并没有针对cuda做优化，所以差距不是很大，可能也和计算图结构关系较大，对于并行结构
cuda+C++的优势就很大了！
JIT编译扩展
注：为了测试这里先卸载llltm_cpp模块
$ pip uninstall lltm_cpp Found existing installation: lltm-cpp 0.0.0 Uninstalling lltm-cpp-0.0.0: Would remove: /opt/conda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg Proceed (y/n)? y Successfully uninstalled lltm-cpp-0.0.0
之前，我提到有两种构建 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=["csrc/lltm.cpp"])
这里我们提供了和setuptools相同的信息。在后台会做如下处理:
- 创建一个临时目录
/tmp/torch_extensions/lltm， - 将Ninja构建文件发送到该临时目录中，
- 将源文件编译到共享库中，
- 将此共享库作为 Python 模块导入。
实际上，如果您传递verbose=True到cpp_extension.load()，您将被告知该过程：
>>> from torch.utils.cpp_extension import load
>>> lltm_cpp = load(name="lltm_cpp", sources=["csrc/lltm.cpp"], verbose=True)
Using /root/.cache/torch_extensions as PyTorch extensions root...
Emitting ninja build file /root/.cache/torch_extensions/lltm_cpp/build.ninja...
Building extension module lltm_cpp...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module lltm_cpp...
>>> lltm_cpp = load(name="lltm_cpp", sources=["csrc/lltm.cpp"], verbose=True)
Using /root/.cache/torch_extensions as PyTorch extensions root...
No modifications detected for re-loaded extension module lltm_cpp, skipping build step...
Loading extension module lltm_cpp...
生成的 Python 模块将与 setuptools 生成的完全相同，但无需维护单独的setup.py构建文件。如果您的构建更复杂，并且确实需要setuptools，那么你可以自己写setup.py文件。第一次运行此行时，需要一些时间，因为扩展程序正在后台编译。由于我们使用 Ninja 构建系统来构建您的源代码，因此重新编译是增量的，因此在您第二次运行 Python 模块时重新加载扩展程序很快，并且如果您不更改扩展程序的源文件，开销也很低。
二、C++与CUDA混合编译
为了真正将我们的实现提升到一个新的水平，我们可以使用自定义 CUDA 内核手动编写前向和后向传递的部分内容。对于 LLTM，这具有特别有效的前景，因为顺序中有大量的逐点操作，它们都可以在单个 CUDA 内核中融合和并行化。让我们看看我们如何编写这样的 CUDA 内核并使用此扩展机制将其与 PyTorch 集成。
注：这里的CUDA编程会比较复杂，需要有GPU架构相关的认知，对于这个简单模型，我们知道kernel-->grid-->block-->thread即可！
构建python调用函数、绑定
编写 CUDA 扩展的一般策略是首先编写一个 C++ 文件，该文件定义将从 Python 调用的函数，并使用 pybind11 将这些函数绑定到 Python。此外，该文件还将声明在 CUDA (.cu ) 文件中定义的函数。然后，C++ 函数将进行一些检查并最终将其调用转发给 CUDA 函数。在 CUDA 文件中，我们编写了实际的 CUDA 内核。然后，该cpp_extension包将负责使用 C++ 的编译器如gcc编译 C++ 源代码，并使用 NVIDIA 的nvcc编译器编译 CUDA 源代码。这确保了每个编译器都处理它最了解编译的文件。最终，它们将被链接到一个共享库中，我们可以从 Python 代码中使用它。
构建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
) {
// shape check
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编程
如您所见，上面python与c++的绑定文件它主要是样板文件，负责检查和转发我们将在 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));
}
上面我们实现了一个简单的sigmoid激活。scalar_t是C++的……，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; // 申请线程数1024，最大不能超过1024
// 申请grid，这里指定了(state_size + threads - 1) / threads行，batch_size列个blocks
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宏和内核启动(由<<<blocks, threads>>>表征)。虽然 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的目的就是处理这个调度。它需要传入三部分内容：
- 一个type：我们这里传入的是
gates.type() - 一个name：用于错误信息打印
- 一个lambda函数
在这个 lambda 函数中，类型别名 scalar_t是可用的，并被定义为张量在运行时在该上下文中实际存在的类型。因此，如果我们有一个模板函数（我们的 CUDA 内核），我们可以用这个scalar_t别名实例化它，然后调用正确的函数。在这种情况下，我们还希望检索张量的数据指针作为该scalar_t类型的指针。如果你想调度所有类型而不仅仅是浮点类型（Float 和Double），你可以使用AT_DISPATCH_ALL_TYPES.
请注意，我们使用纯 ATen 执行一些操作。只使用 ATen 的默认实现，这些操作仍将在 GPU 上运行。这是有道理的，因为 ATen 将使用高度优化的例程来处理矩阵乘法（例如addmm）或卷积，这将更难实现和改进。
至于内核启动(kernel)本身，我们在这里指定每个 CUDA 块将有 1024 个线程，并且整个 GPU Grid被分成尽可能多的线程块，以便每个组件使用一个线程填充我们的矩阵。例如，如果我们的状态大小是 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 内核中使用高级类型不可知的张量是非常低效的。
然而，这是以易用性和可读性为代价的，尤其是对于高维数据。在我们的示例中，很多数据都是dims=3，我们如何访问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() and .stride()帮助我们实现多维索引！.accessor<> 接口旨在有效地访问 cpu 上张量的数据；packed_accessor64<>和 packed_accessor32<>，它产生具有 64 位或 32 位整数索引的cuda张量访问器。
与 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) 会使内核变慢。
CUDA函数声明
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];
}
}
该实现更具可读性！
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操作与python集成
预编译
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='lltm',
ext_modules=[
CUDAExtension('lltm_cuda', [
'csrc/lltm_cuda.cpp',
'csrc/lltm_cuda_kernel.cu',
])
],
cmdclass={
'build_ext': BuildExtension
})
我们现在使用CUDAExtension()，而不是CppExtension(). 我们可以只指定.cu文件和.cpp文件——库会为您处理所有麻烦。
即时编译
from torch.utils.cpp_extension import load
lltm = load(name='lltm', sources=['csrc/lltm_cuda.cpp', 'csrc/lltm_cuda_kernel.cu'])
性能比较
我们希望将代码的逐点操作与 CUDA 并行化和融合可以提高 LLTM 的性能。让我们看看这是否成立。我们可以运行我之前列出的代码来运行基准测试。
$ python test_gpu.py
rnn1 Forward: 17.209 s | Backward 26.507 s (PY)
rnn2 Forward: 15.729 s | Backward 53.925 s (C)
rnn3 Forward: 13.142 s | Backward 25.142 s (CUDA)
关于C在GPU中反向传播的用时异常大，这个有可能是由于容器的资源配置造成的！具体是什么问题，这里先不追究！
注：TODO：
- 反向传播用时异常调试；
- cuda编程资源释放
完！
