Pytorch自定義C++/CUDA擴展

翻譯自：官方文檔

PyTorch 提供了大量與神經網絡、張量代數、數據整理和其他操作。但是，我們有時會需要更加定制化的操作。例如，想要使用論文中找到的一種新型的激活函數，或者實現自己設計的算子。

在 PyTorch 中集成此類自定義操作的最簡單方法是使用 Python 擴展這里概述的 Function 和 Module 。這里已經提供了自動微分的全部功能（無需編寫計算微分的函數）以及 Python 的常用的表達是。但是，有時算子更適合用 C++ 實現。例如，某些算子可能需要非常快，因為它在模型中被非常頻繁地調用，或者即使很少調用也非常耗時。另一個可能的原因是某些算子依賴于其他 C 或 C++ 庫或。為了解決這種情況，PyTorch 提供了一種編寫自定義 C++ 擴展的非常簡單的方法。

C++ 擴展允許用戶創建在 out-of-source 定義的 PyTorch 算子，即與 PyTorch 后端分離。這種方法不同于實現原生 PyTorch 操作的方式。 C++ 擴展旨在節省大量與將操作與 PyTorch 的后端集成相關的樣板，同時為基于 PyTorch 的項目提供高度的靈活性。然而，一旦將算子定義為 C++ 擴展，將其轉換為原生 PyTorch 函數主要是代碼組織問題，如果決定將操作貢獻給上游，則可以事后解決。

動機、實例與Python擴展實現

本文下面的部分將介紹一個編寫和使用 C++（和 CUDA）擴展的實例。

假設我們提出了一種新的循環單元，這個循環單元類似于 LSTM，但不同之處在于它沒有遺忘門，并使用指數線性單元 (ELU) 作為其內部激活函數。因為這個單元永遠不會忘記，我們稱之為 LLTM，或 Long-Long-Term-Memory 單元。

LLTM 與普通 LSTM 的兩種不同之處非常重要，以至于我們無法為我們的目的配置 PyTorch 的 LSTM 算子，因此我們必須創建一個自定義算子。第一個也是最簡單的方法——可能在所有情況下都是很好的第一步——是用 Python 在普通的 PyTorch 中實現我們想要的功能。為此，我們需要繼承 torch.nn.Module 并實現 LLTM 的 forward。即：

import torch
import math
import torch.nn.functional as Fclass LLTM(torch.nn.Module):def __init__(self, input_features, state_size):super(LLTM, self).__init__()self.input_features = input_featuresself.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 = stateX = 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_gatereturn new_h, new_cell

然后我們可以這樣調用：

import torchbatch_size = 4
input_features = 12
state_size = 8X = 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 的融合版本。我們將從使用純 C++ 編寫它開始，使用為 PyTorch 的大部分后端提供支持的 ATen 庫。然后，我們將通過將模型的一部分移動到 CUDA 內核以從 GPU 提供的大規模并行性中受益，從而進一步加快速度。

實現C++擴展

C++ 擴展有兩種形式：它們可以使用

setuptools “ahead of time (AOT)” 構建，
或者通過 torch.utils.cpp_extension.load() “just in time (JIT)” 構建。

我們將逐個介紹。

aot編譯擴展

對于“ahead of time”方式，我們通過編寫一個 setup.py 腳本來構建我們的 C++ 擴展，該腳本使用 setuptools 來編譯我們的 C++ 代碼。

from setuptools import setup, Extension
from torch.utils import cpp_extensionsetup(name='lltm_cpp',ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['lltm.cpp'])],cmdclass={'build_ext': cpp_extension.BuildExtension})

在此代碼中，CppExtension 是 setuptools.Extension 的一個方便的 wrapper，它傳遞正確的包含路徑并將擴展的語言設置為 C++。等效的原 setuptools 代碼是：

Extension(name='lltm_cpp',sources=['lltm.cpp'],include_dirs=cpp_extension.include_paths(),language='c++')

BuildExtension 執行許多必需的配置步驟和檢查，并在混合 C++/CUDA 擴展的情況下管理混合編譯。這就是我們現在真正需要了解的關于構建 C++ 擴展的全部內容！現在讓我們看看我們的 C++ 擴展的實現，即 lltm.cpp。

編寫C++算子

接下來我們開始用 C++ 實現 LLTM。反向傳播需要的一個函數是 sigmoid 的導數。下面一小段代碼，我們據此來討論在編寫 C++ 擴展時的環境：

#include <torch/extension.h>
#include <iostream>torch::Tensor d_sigmoid(torch::Tensor z) {auto s = torch::sigmoid(z);return (1 - s) * s;
}

<torch/extension.h> 是 “一站式” 頭文件，包含編寫 C++ 擴展所需的所有 PyTorch 內容。這包括：

ATen 庫，這是我們用于張量計算的主要 API，
pybind11，這是我們為 C++ 代碼創建 Python 綁定的方式，
以及管理 ATen 和 pybind11 之間交互細節的頭文件。

d_sigmoid() 的實現展示了如何使用 ATen API。 PyTorch 的張量和變量接口是從 ATen 庫自動生成的，因此我們可以或多或少地將 Python 實現 1:1 轉換為 C++。我們所有計算的主要數據類型是 torch::Tensor。可以在這里查看其完整的 API。另請注意，我們可以包含 <iostream> 等 C 或 C++ 頭文件，并支持 C++11 的全部功能。

請注意，在 Windows 上解析 torch/extension.h 時，CUDA-11.5 nvcc 會遇到內部編譯器錯誤。要解決此問題，請將 python 綁定邏輯移動到純 C++ 文件。

示例如下，使用：

#include <ATen/ATen.h>
at::Tensor SigmoidAlphaBlendForwardCuda(....)

而不要：

#include <torch/extension.h>
torch::Tensor SigmoidAlphaBlendForwardCuda(...)

這時 nvcc 的一個bug，目前仍是 open 的 issue，完整解決的代碼示例在這里。

前向傳播

下面，我們給出完整前向傳播的 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 目前沒有為我們提供自動生成反向傳播函數的方法（之前提到 Python 可以）。因此，我們還必須自己實現 LLTM 的反向傳播，它計算損失關于前向傳播的每個輸入的導數。最終，我們將前向和后向函數都放入 torch.autograd.Function 中，來創建一個的 Python binding。反向函數稍微復雜一些，因此我們不會深入研究代碼（如果有興趣，可以閱讀 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.cppsetup.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.pthInstalled ~/local/miniconda/lib/python3.7/site-packages/lltm_cpp-0.0.0-py3.7-linux-x86_64.egg
Processing dependencies for lltm-cpp==0.0.0
Finished processing dependencies for lltm-cpp==0.0.0

關于編譯器的一點說明：由于 ABI 版本控制問題，用于構建 C++ 擴展的編譯器必須與構建 PyTorch 的編譯器 ABI 兼容。實際上，這意味著必須在 Linux 上使用 GCC 4.9 及更高版本。對于 Ubuntu 16.04 和其他較新的 Linux 發行版，這應該已經是默認編譯器了。在 MacOS 上，必須使用 clang（沒有任何 ABI 版本問題）。最麻煩的情況，需要從源代碼構建 PyTorch，然后使用相同的編譯器構建擴展。

構建擴展后，我們可以使用在 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 instanceforward(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_cppclass LLTMFunction(torch.autograd.Function):@staticmethoddef 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@staticmethoddef 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 = outputsreturn d_input, d_weights, d_bias, d_old_h, d_old_cellclass LLTM(torch.nn.Module):def __init__(self, input_features, state_size):super(LLTM, self).__init__()self.input_features = input_featuresself.state_size = state_sizeself.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 timeimport torchbatch_size = 16
input_features = 32
state_size = 128X = 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() - startstart = time.time()(new_h.sum() + new_C.sum()).backward()backward += time.time() - startprint('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

GPU上的性能

關于 PyTorch 的 ATen 后端的一個奇妙的事實是它抽象了我們正在運行的計算設備。這意味著我們為 CPU 編寫的相同代碼也可以在 GPU 上運行，并且各個操作將相應地分派給 GPU 優化的實現。對于矩陣乘法（如 mm 或 addmm）等某些操作，這是一個巨大的勝利。讓我們看看使用 CUDA 張量運行 C++ 代碼可以獲得多少性能。不需要對我們的實現進行任何更改，我們只需要將我們的張量從 Python 放入 GPU 內存中，在創建時添加 device=cuda_device 參數或在創建后使用 .to(cuda_device) ：

import torchassert torch.cuda.is_available()
cuda_device = torch.device("cuda")  # device object representing GPUbatch_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() - startstart = time.time()(new_h.sum() + new_C.sum()).backward()torch.cuda.synchronize()backward += time.time() - startprint('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++ 擴展的方法。

JIT編譯擴展

之前，我提到有兩種構建 C++ 擴展的方法：使用 setuptools 或 just in time (JIT)。講完了前者，我們再來談談后者。 JIT 編譯機制通過調用 PyTorch API 中名為 torch.utils.cpp_extension.load() 的函數，為我們提供了一種即時編譯和加載擴展的方法。對于 LLTM：

from torch.utils.cpp_extension import loadlltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

在這里，我們為該函數提供與 setuptools 相同的信息。在后臺，這將執行以下操作：

創建一個臨時目錄/tmp/torch_extensions/lltm，
將 Ninja 構建文件發送到該臨時目錄中，
將源文件編譯到共享庫中，
將此共享庫作為 Python 模塊導入。

如果想看到這個詳細的過程，可以給 cpp_extension.load() 傳遞參數 verbose=True ：

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 declarationsstd::vector<torch::Tensor> lltm_cuda_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell);std::vector<torch::Tensor> lltm_cuda_backward(torch::Tensor grad_h,torch::Tensor grad_cell,torch::Tensor new_cell,torch::Tensor input_gate,torch::Tensor output_gate,torch::Tensor candidate_cell,torch::Tensor X,torch::Tensor gate_weights,torch::Tensor weights);// C++ interface#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)std::vector<torch::Tensor> lltm_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell) {CHECK_INPUT(input);CHECK_INPUT(weights);CHECK_INPUT(bias);CHECK_INPUT(old_h);CHECK_INPUT(old_cell);return lltm_cuda_forward(input, weights, bias, old_h, old_cell);
}std::vector<torch::Tensor> lltm_backward(torch::Tensor grad_h,torch::Tensor grad_cell,torch::Tensor new_cell,torch::Tensor input_gate,torch::Tensor output_gate,torch::Tensor candidate_cell,torch::Tensor X,torch::Tensor gate_weights,torch::Tensor weights) {CHECK_INPUT(grad_h);CHECK_INPUT(grad_cell);CHECK_INPUT(input_gate);CHECK_INPUT(output_gate);CHECK_INPUT(candidate_cell);CHECK_INPUT(X);CHECK_INPUT(gate_weights);CHECK_INPUT(weights);return lltm_cuda_backward(grad_h,grad_cell,new_cell,input_gate,output_gate,candidate_cell,X,gate_weights,weights);
}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &lltm_forward, "LLTM forward (CUDA)");m.def("backward", &lltm_backward, "LLTM backward (CUDA)");
}

如我們所見，它主要是樣板，檢查和轉發到我們將在 CUDA 文件中定義的函數。我們將此文件命名為 lltm_cuda_kernel.cu（注意 .cu 擴展名！）。 NVCC 可以編譯 C++11，因此我們仍然可以使用 ATen 和 C++ 標準庫（但不是 torch.h）。請注意，setuptools 無法處理具有相同名稱但擴展名不同的文件，因此如果使用 setup.py 方法而不是 JIT 方法，則必須為 CUDA 文件指定一個與 C++ 文件不同的名稱（對于 JIT 方法，lltm. cpp 和 lltm.cu 可以正常工作）。讓我們看一下這個文件的樣子：

#include <torch/extension.h>#include <cuda.h>
#include <cuda_runtime.h>#include <vector>template <typename scalar_t>
__device__ __forceinline__ scalar_t sigmoid(scalar_t z) {return 1.0 / (1.0 + exp(-z));
}

在這里，我們看到了我剛剛描述的頭文件，以及我們使用 CUDA 特定語法（如 __device__ 和 __forceinline__ ）和函數（如 exp）。我們繼續使用一些輔助函數：

template <typename scalar_t>
__device__ __forceinline__ scalar_t d_sigmoid(scalar_t z) {const auto s = sigmoid(z);return (1.0 - s) * s;
}template <typename scalar_t>
__device__ __forceinline__ scalar_t d_tanh(scalar_t z) {const auto t = tanh(z);return 1 - (t * t);
}template <typename scalar_t>
__device__ __forceinline__ scalar_t elu(scalar_t z, scalar_t alpha = 1.0) {return fmax(0.0, z) + fmin(0.0, alpha * (exp(z) - 1.0));
}template <typename scalar_t>
__device__ __forceinline__ scalar_t d_elu(scalar_t z, scalar_t alpha = 1.0) {const auto e = exp(z);const auto d_relu = z < 0.0 ? 0.0 : 1.0;return d_relu + (((alpha * (e - 1.0)) < 0.0) ? (alpha * e) : 0.0);
}

現在要實際實現一個函數，我們還需要兩件事：一個執行我們不希望顯式手動編寫的操作并調用 CUDA 內核的函數，然后是我們想要加速的部分的實際 CUDA 內核 . 對于前向傳遞，第一個函數應該如下所示：

std::vector<torch::Tensor> lltm_cuda_forward(torch::Tensor input,torch::Tensor weights,torch::Tensor bias,torch::Tensor old_h,torch::Tensor old_cell) {auto X = torch::cat({old_h, input}, /*dim=*/1);auto gates = torch::addmm(bias, X, weights.transpose(0, 1));const auto batch_size = old_cell.size(0);const auto state_size = old_cell.size(1);auto new_h = torch::zeros_like(old_cell);auto new_cell = torch::zeros_like(old_cell);auto input_gate = torch::zeros_like(old_cell);auto output_gate = torch::zeros_like(old_cell);auto candidate_cell = torch::zeros_like(old_cell);const int threads = 1024;const dim3 blocks((state_size + threads - 1) / threads, batch_size);AT_DISPATCH_FLOATING_TYPES(gates.type(), "lltm_forward_cuda", ([&] {lltm_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(gates.data<scalar_t>(),old_cell.data<scalar_t>(),new_h.data<scalar_t>(),new_cell.data<scalar_t>(),input_gate.data<scalar_t>(),output_gate.data<scalar_t>(),candidate_cell.data<scalar_t>(),state_size);}));return {new_h, new_cell, input_gate, output_gate, candidate_cell, X, gates};
}

這里的關鍵是 AT_DISPATCH_FLOATING_TYPES 宏和內核啟動（由 <<<...>>> 指示）。雖然 ATen 抽象出我們處理的張量的設備和數據類型，但張量在運行時仍將由具體設備上的具體類型的內存支持。因此，我們需要一種在運行時確定張量是什么類型的方法，然后選擇性地調用具有相應正確類型簽名的函數。手動完成，如下：

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

AT_DISPATCH_FLOATING_TYPES 的目的是為我們處理這個調度。它需要一個類型（在我們的例子中是 gates.type()）、一個名稱（用于顯示錯誤消息）和一個 lambda 函數。在這個 lambda 函數中，類型別名 scalar_t 可用，并被定義為張量在運行時在該 context 中實際存在的類型。因此，如果我們有一個模板函數（就是我們的 CUDA 核函數），我們可以用這個 scalar_t 別名實例化它，然后調用正確的函數。在這種情況下，我們還希望檢索張量的數據指針作為該 scalar_t 類型的指針。如果想分派所有類型而不僅僅是浮點類型（Float 和 Double），可以使用 AT_DISPATCH_ALL_TYPES。

請注意，我們使用純 ATen 執行一些操作。這些操作仍將在 GPU 上運行，但使用 ATen 的默認實現。這是有道理的，因為 ATen 將使用高度優化的例程來處理矩陣乘法（例如 addmm）或卷積，這部分很難實現和改進。

至于內核啟動本身，我們在此指定每個 CUDA block 有 1024 個線程，并且整個 GPU 網格被分成多個 1 x 1024 線程塊，以每個組件一個線程填充我們的矩陣。例如，如果我們的 state_size 大小為 2048，batch_size 大小為 4，我們將啟動總共 4 x 2 = 8 個塊，每個 1024 個線程。如果不了解 CUDA“block”或“grid”，那么可以參考閱讀有關 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 循環超過一百萬個串行元素來做到這一點，就會明白為什么這會快得多。

使用accessors

可以在 CUDA 核中看到我們直接處理對應類型的指針。事實上，直接在 cuda 核中使用高級類型不可知的張量是非常低效的。然而，這是以易用性和可讀性為代價的，尤其是對于高維數據。在我們的示例中，我們知道例如連續門張量具有 3 個維度：

batch，大小為 batch_size，步長為 3*state_size
row，大小為 3，步長為 state_size
index，大小為 state_size ，步長為 1

那么我們如何訪問內核中的元素 gates[n][row][column] 呢？事實證明，可以通過一些簡單的算術來訪問對應位置的元素。

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

這個表達式需要明確知道步長，并在其參數中傳遞給核函數。可以看到，在核函數接受多個不同大小的張量的情況下，最終會得到一個很長的參數列表。對我們來說幸運的是，ATen 提供了通過單一動態檢查創建的訪問器，張量是維度的類型和數量。然后訪問器公開一個 API 以有效地訪問張量元素，而無需轉換為單個指針：

torch::Tensor foo = torch::rand({12, 12});// assert foo is 2-dimensional and holds floats.
auto foo_a = foo.accessor<float,2>();
float trace = 0;for(int i = 0; i < foo_a.size(0); i++) {// use the accessor foo_a to get tensor data.trace += foo_a[i][i];
}

Accessor 對象有一個相對高層的接口，具有 .size() 和 .stride() 方法以及多維索引。 .accessor<> 接口旨在有效地訪問 cpu 張量上的數據。對應于 cuda 張量的是 packed_accessor64<> 和 packed_accessor32<>，它們生成具有 64 位或 32 位整數索引的打包 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 與常規 accessor 相同。參數 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 indexconst int n = blockIdx.y;// column indexconst 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];}
}

該實現更具可讀性。然后通過在 host 函數中使用 .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 indexconst int n = blockIdx.y;// column indexconst 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, CUDAExtensionsetup(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 loadlltm = 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++ 擴展機制以及使用它們的動機有了很好的了解。可以在這里找到本文中展示的代碼示例。如果有任何問題，請使用論壇，或者看常見問題的 FAQ。