PyTorch中构建和调用C++/CUDA扩展

最新推荐文章于 2024-06-15 07:51:23 发布

NaiveYoungPeo

最新推荐文章于 2024-06-15 07:51:23 发布

阅读量3.7k

点赞数 6

分类专栏： PyTorch 文章标签： pytorch c++ python

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

版权

PyTorch 专栏收录该内容

5 篇文章 1 订阅

订阅专栏

翻译至PyTorch官方教程
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++。等效的 vanilla 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。另请注意，我们可以包含或任何其他 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。 torch扩展构建将其定义为您在 setup.py 脚本中为扩展指定的名称。在这种情况下，TORCH_EXTENSION_NAME 的值将是“lltm_cpp”。这是为了避免必须在两个地方（构建脚本和 C++ 代码）维护扩展名，因为两者之间的不匹配会导致令人讨厌且难以跟踪的问题。

使用您的扩展

我们现在准备在 PyTorch 中导入我们的扩展。此时，您的目录结构可能如下所示：

pytorch/
    lltm-extension/
    lltm.cpp
    setup.py

现在，运行 python setup.py install 来构建和安装你的扩展。

关于编译器的一个小说明：由于 ABI 版本问题，用于构建 C++ 扩展的编译器必须与构建 PyTorch 的编译器 ABI 兼容。实际上，这意味着您必须在 Linux 上使用 GCC 4.9 及更高版本。对于 Ubuntu 16.04 和其他更新的 Linux 发行版，这应该已经是默认编译器。在 MacOS 上，您必须使用 clang（它没有任何 ABI 版本控制问题）。在最坏的情况下，您可以使用编译器从源代码构建 PyTorch，然后使用相同的编译器构建扩展。

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

In [1]: import torch
In [2]: import lltm_cpp
In [3]: lltm_cpp.forward
Out[3]: <function lltm.PyCapsule.forward>

如果我们在函数或模块上调用 help()，我们可以看到它的签名与我们的 C++ 代码匹配：

In[4] help(lltm_cpp.forward)
forward(...) method of builtins.PyCapsule 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} 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 或 just in time (JIT)。讲完前者，我们来详细介绍后者。 JIT 编译机制通过调用 PyTorch 的 API 中名为 torch.utils.cpp_extension.load() 的简单函数，为您提供了一种动态编译和加载扩展的方法。对于 LLTM，这看起来很简单：

from torch.utils.cpp_extension import load

lltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

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

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

实际上，如果您将verbose=True 传递给cpp_extension.load()，您将被告知该过程：

Using /tmp/torch_extensions as PyTorch extensions root...
Emitting ninja build file /tmp/torch_extensions/lltm_cpp/build.ninja...
Building extension module lltm_cpp...
Loading extension module lltm_cpp...

生成的 Python 模块将与 setuptools 生成的完全相同，但不需要维护单独的 setup.py 构建文件。如果您的设置更复杂，并且您确实需要 setuptools 的全部功能，您可以编写自己的 setup.py，但在许多情况下，这种 JIT 技术会很好。第一次运行此行时，需要一些时间，因为扩展程序正在后台编译。由于我们使用 Ninja 构建系统来构建您的源代码，重新编译是增量的，因此如果您不更改扩展的源文件，第二次运行 Python 模块时重新加载扩展会很快并且开销很低。

编写混合 C++/CUDA 扩展

为了真正将我们的实现提升到一个新的水平，我们可以使用自定义 CUDA 内核手写部分向前和向后传递。对于 LLTM，这有可能特别有效，因为有大量按顺序进行的逐点操作，它们都可以在单个 CUDA 内核中融合和并行化。让我们看看如何编写这样的 CUDA 内核并使用这种扩展机制将其与 PyTorch 集成。

编写 CUDA 扩展的一般策略是首先编写一个 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 循环超过一百万个连续元素来做到这一点，你就会明白为什么这会快得多。

使用Accessor

您可以在 CUDA 内核中看到我们直接处理具有正确类型的指针。事实上，直接在 CUDA 内核中使用高级类型不可知张量是非常低效的。

然而，这是以易用性和可读性为代价的，尤其是对于高维数据。在我们的例子中，我们知道例如连续门张量有 3 个维度：

batch, size of batch_size and stride of 3*state_size
row, size of 3 and stride of state_size
index, size of state_size and stride of 1

那么我们如何访问内核中的元素 gates[n][row][column] 呢？事实证明，您需要使用一些简单的算法来访问元素的步幅。

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

除了冗长之外，该表达式还需要 stride 才能明确知道，从而在其参数中传递给内核函数。你可以看到，在核函数接受多个不同大小的张量的情况下，你最终会得到一个很长的参数列表。

对我们来说幸运的是，ATen 提供了通过单个动态检查创建的Accessor，即张量是维度的类型和数量。然后，访问器公开一个 API，用于高效访问 Tensor 元素，而无需转换为单个指针：

torch::Tensor foo = torch::rand({12, 12});

// assert foo is 2-dimensional and holds floats.
auto foo_a = foo.accessor<float,2>();
float trace = 0;

for(int i = 0; i < foo_a.size(0); i++)
{
    // use the accessor foo_a to get tensor data.
    trace += foo_a[i][i];
}

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

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

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

__global__ void lltm_cuda_forward_kernel(
    const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> gates,
    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> old_cell,
    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> new_h,
    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> new_cell,
    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> input_gate,
    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> output_gate,
    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> candidate_cell)

让我们分解这里使用的模板。前两个参数 scalar_t 和 2 与常规访问器相同。参数 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
    })

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