1、前言?

?一般情況下，pytorch推薦使用python層的前端語言來構建新的算子。因為pytorch在python層的api已經足夠豐富，可以構造出很多自定義的算子。但是有時候出于一些其他方面的考慮，會需要增加底層算子。例如有時候對性能要求很高，python不滿足需求，又或者是需要鏈接其他的動態庫（blas，mkl等），因此pytorch也提供了直接擴展底層C++算子的能力。主要有三種方式，native_functions.yaml、C++ extension方式、?OP register?方式。

2、native_functions.yaml方式

pytorch的原生算子很多都是使用這種方式組織的。在native_functions.yaml中有關于各個算子的說明，然后在同級目錄下面有這些算子的實現。使用該方式添加新的算子，主要用在已經支持的硬件上面。例如pytorch本身已經支持了CPU和GPU，此時需要一些新的算子，該算子只需要在CPU或者GPU上面運行，那么這種方式就非常適合。只需要定義新算子的kernel實現，然后添加配置信息，就可以自動生成：torch.xxx()、torch.nn.functional.xxx()以及tensor.xxx()方法，而不用去關注算子與pytorch是如何銜接，以及如何把算子添加到tensor的屬性中等其他細節。native_functions.yaml文件位于pytorch源碼的pytorch/aten/src/Aten/native/native_functions.yaml，內容如下（截取absolute算子的配置信息），對于每個算子的描述，包括幾個主要字段：func、variants、dispatch等。

func字段：表示算子的名稱以及輸入輸出參數類型

variants字段：表示需要自動生成的高級方法。function表示自動生成torch.absolute()方法，method表示生成?tensor的absolute ()方法，即可以定義一個tensor a，然后可以執行a.absolute()方法。

dispatch字段：表示分發的設備類型對應的op方法。CPU指的是該算子支持CPU設備，對應的實現函數為abs函數，CUDA指的是當前算子支持GPU設備，對應的實現函數為cuda的abs函數。

下面以pytorch自帶的leakly_relu算子來具體分析添加算子的流程。首先是需要在native_functions.yaml中添加算子的說明，包括反向傳播函數。如下代碼片段中的leaky_relu和leaky_relu_backward函數說明。這里的python_module:nn，表示將該方法自動生成到torch.nn.functional模塊中，這樣就可以通過torch.nn.functional.leaky_relu來調用這個算子。

- func: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensoruse_c10_dispatcher: fullpython_module: nndispatch:CPU: leaky_reluCUDA: leaky_reluQuantizedCPU: quantized_leaky_relu- func: leaky_relu_backward(Tensor grad_output, Tensor self, Scalar negative_slope, bool self_is_result) -> Tensoruse_c10_dispatcher: fullpython_module: nn

其次，需要在配置文件tools/autograd/derivatives.yaml中添加算子和反向算子的對應關系，如下代碼段表示，即說明了leaky_relu的反向傳播函數為leaky_relu_backward。

- name: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensorself: leaky_relu_backward(grad, self, negative_slope, false)

完成了算子的說明之后，需要在aten/src/Aten/native/目錄下面通過C++實現相關的算子流程。Pytorch原生的算子一般按照功能實現在一起。例如激活函數都放在Activation.h與Activation.cpp中。所以leaky_relu的實現就在aten/src/Aten/native/目錄下的Activation.h與Activation.cpp文件中。如下代碼段所示。不過這里定義的實現只是一個封裝，沒有真正的實現。leak_relu調用了leaky_relu_stub方法，leak_relu_backward調用了leak_relu_backward_stub方法。

//Activation.h頭文件 using leaky_relu_fn = void (*)(TensorIterator&, Scalar); using leaky_relu_backward_fn = void (*)(TensorIterator&, Scalar); DECLARE_DISPATCH(leaky_relu_fn, leaky_relu_stub); DECLARE_DISPATCH(leaky_relu_backward_fn, leaky_relu_backward_stub);//Activation.cpp文件內容 DEFINE_DISPATCH(leaky_relu_stub); DEFINE_DISPATCH(leaky_relu_backward_stub);Tensor leaky_relu(const Tensor& self,Scalar negval) {Tensor result;auto iter = TensorIterator::unary_op(result, self);leaky_relu_stub(iter.device_type(), iter, negval);return iter.output(); }Tensor leaky_relu_backward(const Tensor& grad_output,const Tensor& self_or_result,Scalar negval,bool is_result) {Tensor result;auto iter = TensorIterator::binary_op(result, self_or_result, grad_output);leaky_relu_backward_stub(iter.device_type(), iter, negval);return iter.output(); }

最終，CPU端的leaky_relu_stub和leak_relu_backward_stub兩個函數的實現流程都在aten/src/Aten/native/cpu/Activation.cpp中。并且增加了兩個DISPATH（函數分發的說明）。如下代碼段所示：

REGISTER_DISPATCH(leaky_relu_stub, &leaky_relu_kernel); REGISTER_DISPATCH(leaky_relu_backward_stub, &leaky_relu_backward_kernel);static void leaky_relu_kernel(TensorIterator& iter, Scalar negval_) {AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "leaky_relu_cpu", [&] {using Vec = Vec256<scalar_t>;auto zero_vec = Vec((scalar_t)(0));auto one_vec = Vec((scalar_t)(1));scalar_t negval = negval_.to<scalar_t>();Vec negval_v = Vec(negval);cpu_kernel_vec(iter,[&](scalar_t a) -> scalar_t {return a > scalar_t(0) ? a : a * negval;},[&](Vec a) -> Vec {auto r = Vec::blendv(negval_v, one_vec, a > zero_vec);return a * r;});}); }static void leaky_relu_backward_kernel(TensorIterator& iter, Scalar negval_) {AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "leaky_relu_backward_cpu", [&] {using Vec = Vec256<scalar_t>;auto zero_vec = Vec((scalar_t)(0));auto one_vec = Vec((scalar_t)(1));scalar_t negval = negval_.to<scalar_t>();Vec negval_v = Vec(negval);cpu_kernel_vec(iter,[&](scalar_t a, scalar_t b) -> scalar_t {return a > scalar_t(0) ? b : b * negval;},[&](Vec a, Vec b) -> Vec {auto r = Vec::blendv(negval_v, one_vec, a > zero_vec);return b * r;});}); }

同樣的，GPU端的leaky_relu_stub和leak_relu_backward_stub兩個函數的實現流程都在aten/src/Aten/native/cuda/Activation.cu中。并且增加了兩個DISPATH（函數分發的說明）。如下代碼段所示：

REGISTER_DISPATCH(leaky_relu_stub, &leaky_relu_kernel); REGISTER_DISPATCH(leaky_relu_backward_stub, &leaky_relu_backward_kernel);void leaky_relu_kernel(TensorIterator& iter, Scalar negval_) {AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "leaky_relu_cuda", [&]() {AT_SKIP_BFLOAT16_IF_NOT_ROCM(scalar_t, "leaky_relu_cuda", [&] {auto negval = negval_.to<scalar_t>();gpu_kernel(iter, [negval]GPU_LAMBDA(scalar_t a) -> scalar_t {return a > scalar_t(0) ? a : a * negval;});});}); }void leaky_relu_backward_kernel(TensorIterator& iter, Scalar negval_) {AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "leaky_relu_backward_cuda", [&]() {AT_SKIP_BFLOAT16_IF_NOT_ROCM(scalar_t, "leaky_relu_backward_cuda", [&] {auto negval = negval_.to<scalar_t>();gpu_kernel(iter, [negval]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {return a > scalar_t(0) ? b : b * negval;});});}); }

至此，就完成了整個leaky_relu算子的實現流程，總體流程還是比較簡單清晰的，并且只需要考慮算子本身的具體實現，而不需要去考慮如何將算子添加到torch模塊，添加到torch.nn.functional模塊，如何與tensor耦合等業務邏輯。下面這個圖更加清晰的展示了這種實現方式（為了節約圖片高度，省略cuda的實現）。

?

下面以實現一個自定義的xxx算子為例，為了簡單起見，只實現該算子的CPU前向算子。首先在native_functions.yaml文件中增加xxx算子的描述：

- func: xxx(Tensor self) -> Tensoruse_c10_dispatcher: fullpython_module: nndispatch:CPU: xxx

然后在同級目錄下實現算子的表層實現文件，同樣為了簡單起見，直接實現在pytorch已有的Activation.h與Activation.cpp源文件中。如下所示：

//Activation.h文件 using xxx_fn = void (*)(TensorIterator&); DECLARE_DISPATCH(xxx_fn, xxx_stub);//Activation.cpp文件 DEFINE_DISPATCH(xxx_stub);Tensor xxx(const Tensor& self) {Tensor result;auto iter = TensorIterator::unary_op(result, self);xxx_stub(iter.device_type(), iter);return iter.output(); }

最后在cpu/Activation.cpp中實現真正的xxx_stub方法。為了簡單，不做任何數值操作，只是調用printf打印相關信息。

REGISTER_DISPATCH(xxx_stub, &xxx_kernel);static void xxx_kernel(TensorIterator& iter) {AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "xxx_cpu", [&] {printf("xxx op forward!");}); }

編譯之后，對xxx算子進行測試，如下所示：

>>> import torch >>> t = torch.ones(1,3,2,2) >>> t = torch.xxx(t) xxx op forward! >>> t = t.xxx() xxx op forward! 成功打印，說明xxx算子已經可以使用

3、C++?extention方式

雖然native_functions.yaml方式可以比較方便的增加或者修改算子，但是存在一個比較嚴重的問題。就是與pytorch的耦合度過高，由于在pytorch的源碼中直接修改，那么每次增加或者修改算子都需要重新編譯pytorch。為此，pytorch提供了另外一種更加簡便的方式來擴展底層算子，就是?C++ extension方式。它與pytorch的相互解耦，分開編譯，所以增加算子不需要修改pytorh的源碼。它的原理其實就是通過pybind11，將C++編譯為pytroch的一個模塊，這樣就可以在pytorch中通過這個新的模塊來執行新的OP了。?這里以一個小例子來說明如何通過C++extension增加一個算子。該例子出自官方文檔：https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-mixed-c-cuda-extension

算子名稱為lltm，首先看一下目錄結構：

在lltm.cpp中編寫前向和反向函數的功能實現：

#include <vector>std::vector<at::Tensor> lltm_forward(…) {……return {…}; }std::vector<torch::Tensor> lltm_backward(…) {……return {…}; }

另外需要增加pybind11的綁定說明。因為pytorch的c++extension是通過pybind11綁定到python的。綁定說明如下：

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

最后編寫setup.py，用于編譯生成相關的模塊

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})

完成上述代碼的編寫之后，執行：python?setup.py?install即可完成編譯。生成pytorch中可用的lltm算子。下面對新增加的lltm算子進行測試，發現pytorch已經可以準確識別該算子了。

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

4、OP Register方式

雖然C++ extension方式能夠比較方便的增加底層算子。但是也存在一點缺陷。首先它是作為一個額外的擴展模塊接入pytorch，所以在調用這些方法的時候，都是需要直接導入方法名稱。即無法通過torch.xxx或者tensor.xxx的方式進行調用，另外只能支持現有平臺，無法擴展到新的硬件平臺。所以Pytorch還提供了一種更加強大的算子擴展能力，就是OP Register（算子注冊）方式。同樣，該方式與pytorch源碼解耦，增加和修改算子不需要重新編譯pytorch源碼。關于該部分的說明，pytroch的官方文檔中并沒有找到相關信息，但是在pytroch源碼的aten/src/ATen/core/op_registration/README.md中有一些介紹。（備注：雖然該方法與pytorch本身解耦，如果需要增加新硬件平臺對應的算子，那么需要首先在pytroch源碼中增加對新硬件的支持，以及算子分發的DISPATH_KEY等相關信息，然后才能使用該方法注冊基于該新硬件的算子）

用該方式注冊一個新的算子，流程非常簡單：先編寫C++相關的算子實現，然后通過pytorch底層的注冊接口（torch::RegisterOperators），將該算子注冊即可。如下代碼段所示。這里只注冊了pytroch原生支持的CPU和CUDA硬件平臺。

//my_kernel 定義（包括CPU和GPU版本） my_namespace { Tensor my_op_cpu(const Tensor& a, const Tensor& b) {...} Tensor my_op_cuda(const Tensor& a, const Tensor& b) {...} }static auto registry = torch::RegisterOperators().op("my_namespace::my_op", torch::RegisterOperators::options().kernel<decltype(my_kernel_cpu), &my_kernel_cpu>(CPU())).op("my_namespace::my_op", torch::RegisterOperators::options().kernel<decltype(my_kernel_cuda), &my_kernel_cuda>(CUDA()));

如果需要增加新硬件平臺的支持，那么首先需要在pytorch源碼中的Backend、Device等模塊中添加新硬件的支持。假設新硬件平臺名為：VD(Virtual Device)，那么注冊基于VD的新算子就是：

static auto registry = torch::RegisterOperators().op("my_namespace::my_op", torch::RegisterOperators::options().kernel<decltype(my_kernel_cpu), &my_kernel_vd>(VD()))

總結

以上是生活随笔為你收集整理的这可能是关于Pytorch底层算子扩展最详细的总结了！的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。