自定义 C++ 和 CUDA 扩展

原文：https://pytorch.org/tutorials/advanced/cpp_extension.html

作者： Peter Goldsborough

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 或长期记忆单元。

LLTM 与普通 LSTM 的两种区别非常重要，以至于我们无法为自己的目的配置 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 库以普通的 C++ 语言编写代码，该库为 PyTorch 的许多后端提供了强大的支持，并了解它如何使我们轻松转换 Python 代码。 然后，我们将模型的某些部分移至 CUDA 内核，以从 GPU 提供的大量并行处理中受益，从而进一步加快处理速度。

编写 C++ 扩展

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

使用setuptools构建

为了“提前”，我们通过编写一个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', ['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++ 扩展的全部信息！ 现在让我们看一下lltm.cpp中 C++ 扩展的实现。

编写 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 的全部功能。

正向传播

接下来，我们可以将整个正向传播到 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。 火炬扩展程序构建会将其定义为您在setup.py脚本中为扩展程序指定的名称。 在这种情况下，TORCH_EXTENSION_NAME的值为lltm。 这是为了避免在两个位置（构建脚本和 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 编译器兼容。 实际上，这意味着您必须在 Linux 上使用 GCC 4.9 及更高版本。 对于 Ubuntu 16.04 和其他较新的 Linux 发行版，这应该已经是默认的编译器。 在 MacOS 上，您必须使用 clang（它没有任何 ABI 版本控制问题）。 在最坏的情况下，您可以使用编译器从源代码构建 PyTorch，然后使用相同的编译器构建扩展。

扩展程序构建完成后，您可以使用setup.py脚本中指定的名称，将其简单地导入 Python。 只需确保先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_variables)
        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} us | Backward {:.3f} us'.format(forward * 1e6/1e5, backward * 1e6/1e5))

如果我们使用本文开头用纯 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))

再次将普通的 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相同的信息。 在后台，这将执行以下操作：

1. 创建一个临时目录/tmp/torch_extensions/lltm，
2. 将 Ninja 构建文件发送到该临时目录中，
3. 将您的源文件编译到共享库中，
4. 将此共享库导入为 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 扩展的一般策略是首先编写一个 C++ 文件，该文件定义将从 Python 调用的函数，然后使用pybind11将这些函数绑定到 Python。 此外，此文件还将声明在 CUDA（.cu）文件中定义的函数。 然后，C++ 函数将进行一些检查，并最终将其调用转发给 CUDA 函数。 在 CUDA 文件中，我们编写了实际的 CUDA 内核。 然后cpp_extension包将负责使用gcc等 C++ 编译器来编译 C++ 源代码，并使用 NVIDIA 的nvcc编译器来编译 CUDA 源。 这样可以确保每个编译器都照顾最了解要编译的文件。 最终，它们将被链接到一个共享库中，该库可从 Python 代码中获得。

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

#include <torch/extension.h>

#include <vector>

// CUDA forward 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.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

std::vector<torch::Tensor> lltm_forward(
    torch::Tensor input,
    torch::Tensor weights,
    torch::Tensor bias,
    torch::Tensor old_h,
    torch::Tensor old_cell) {
  CHECK_INPUT(input);
  CHECK_INPUT(weights);
  CHECK_INPUT(bias);
  CHECK_INPUT(old_h);
  CHECK_INPUT(old_cell);

  return lltm_cuda_forward(input, weights, bias, old_h, old_cell);
}

std::vector<torch::Tensor> lltm_backward(
    torch::Tensor grad_h,
    torch::Tensor grad_cell,
    torch::Tensor new_cell,
    torch::Tensor input_gate,
    torch::Tensor output_gate,
    torch::Tensor candidate_cell,
    torch::Tensor X,
    torch::Tensor gate_weights,
    torch::Tensor weights) {
  CHECK_INPUT(grad_h);
  CHECK_INPUT(grad_cell);
  CHECK_INPUT(input_gate);
  CHECK_INPUT(output_gate);
  CHECK_INPUT(candidate_cell);
  CHECK_INPUT(X);
  CHECK_INPUT(gate_weights);
  CHECK_INPUT(weights);

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

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

如您所见，它主要是样板文件，检查并转发到我们将在 CUDA 文件中定义的功能。 我们将此文件命名为lltm_cuda_kernel.cu（请注意.cu扩展名！）。 NVCC 可以合理地编译 C++ 11，因此我们仍然可以使用 ATen 和 C++ 标准库（但不能使用torch.h）。 请注意，setuptools无法处理具有相同名称但扩展名不同的文件，因此，如果您使用setup.py方法而不是 JIT 方法，则必须给 CUDA 文件指定一个与 C++ 文件不同的名称（对于 JIT 方法， lltm.cpp和lltm.cu可以正常工作）。 让我们看一下该文件的外观：

#include <torch/extension.h>

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

#include <vector>

template <typename scalar_t>
__device__ __forceinline__ scalar_t sigmoid(scalar_t z) {
  return 1.0 / (1.0 + exp(-z));
}

在这里，我们看到了我刚刚描述的标头，以及我们正在使用特定于 CUDA 的声明，例如__device__和__forceinline__以及类似exp的事实。 让我们继续一些我们需要的辅助功​​能：

template <typename scalar_t>
__device__ __forceinline__ scalar_t d_sigmoid(scalar_t z) {
  const auto s = sigmoid(z);
  return (1.0 - s) * s;
}

template <typename scalar_t>
__device__ __forceinline__ scalar_t d_tanh(scalar_t z) {
  const auto t = tanh(z);
  return 1 - (t * t);
}

template <typename scalar_t>
__device__ __forceinline__ scalar_t elu(scalar_t z, scalar_t alpha = 1.0) {
  return fmax(0.0, z) + fmin(0.0, alpha * (exp(z) - 1.0));
}

template <typename scalar_t>
__device__ __forceinline__ scalar_t d_elu(scalar_t z, scalar_t alpha = 1.0) {
  const auto e = exp(z);
  const auto d_relu = z < 0.0 ? 0.0 : 1.0;
  return d_relu + (((alpha * (e - 1.0)) < 0.0) ? (alpha * e) : 0.0);
}

现在，要真正实现一个函数，我们再次需要两件事：一个函数执行我们不想手工明确编写的操作并调用 CUDA 内核，然后是要加速的部分的实际 CUDA 内核。 。 对于正向传播，第一个函数应如下所示：

std::vector<torch::Tensor> lltm_cuda_forward(
    torch::Tensor input,
    torch::Tensor weights,
    torch::Tensor bias,
    torch::Tensor old_h,
    torch::Tensor old_cell) {
  auto X = torch::cat({old_h, input}, /*dim=*/1);
  auto gates = torch::addmm(bias, X, weights.transpose(0, 1));

  const auto batch_size = old_cell.size(0);
  const auto state_size = old_cell.size(1);

  auto new_h = torch::zeros_like(old_cell);
  auto new_cell = torch::zeros_like(old_cell);
  auto input_gate = torch::zeros_like(old_cell);
  auto output_gate = torch::zeros_like(old_cell);
  auto candidate_cell = torch::zeros_like(old_cell);

  const int threads = 1024;
  const dim3 blocks((state_size + threads - 1) / threads, batch_size);

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

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

这里的主要关注点是AT_DISPATCH_FLOATING_TYPES宏和内核启动（由<<<...>>>指示）。 尽管 ATen 提取了我们处理过的张量的设备和数据类型，但张量在运行时仍将由具体设备上具体类型的内存支持。 因此，我们需要一种在运行时确定张量是什么类型，然后有选择地调用具有相应正确类型签名的函数的方法。 手动完成后，（在概念上）将如下所示：

switch (tensor.type().scalarType()) {
  case torch::ScalarType::Double:
    return function<double>(tensor.data<double>());
  case torch::ScalarType::Float:
    return function<float>(tensor.data<float>());
  ...
}

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

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

至于内核启动本身，我们在这里指定每个 CUDA 块将具有 1024 个线程，并且将整个 GPU 网格分为所需的1 x 1024线程块，以便用每个组件一个线程填充矩阵。 例如，如果我们的状态大小为 2048，批量大小为 4，则我们将以每个 1024 个线程总共启动4 x 2 = 8块。 如果您以前从未听说过 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 个维度：

1. 批量，batch_size的大小和3*state_size的步幅
2. 3的行，大小和state_size的步幅
3. 指数，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()方法以及多维索引。 .accessor<>接口旨在在 CPU 张量上有效访问数据。 CUDA 张量的等效项是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_forward_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
    })

现在，我们使用CUDAExtension()代替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++ 扩展机制有了一个很好的了解，并有使用它们的动机。 您可以在此处找到本说明中显示的代码示例。 如有疑问，请使用论坛。 另外，请务必查看我们的常见问题解答，以防遇到任何问题。