1、前言?

?一般情況下，pytorch推薦使用python層的前端語言來構(gòu)建新的算子。因?yàn)閜ytorch在python層的api已經(jīng)足夠豐富，可以構(gòu)造出很多自定義的算子。但是有時(shí)候出于一些其他方面的考慮，會(huì)需要增加底層算子。例如有時(shí)候?qū)π阅芤蠛芨?#xff0c;python不滿足需求，又或者是需要鏈接其他的動(dòng)態(tài)庫（blas，mkl等），因此pytorch也提供了直接擴(kuò)展底層C++算子的能力。主要有三種方式，native_functions.yaml、C++ extension方式、?OP register?方式。

2、native_functions.yaml方式

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

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

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

dispatch字段：表示分發(fā)的設(shè)備類型對(duì)應(yīng)的op方法。CPU指的是該算子支持CPU設(shè)備，對(duì)應(yīng)的實(shí)現(xiàn)函數(shù)為abs函數(shù)，CUDA指的是當(dāng)前算子支持GPU設(shè)備，對(duì)應(yīng)的實(shí)現(xiàn)函數(shù)為cuda的abs函數(shù)。

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

- 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中添加算子和反向算子的對(duì)應(yīng)關(guān)系，如下代碼段表示，即說明了leaky_relu的反向傳播函數(shù)為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++實(shí)現(xiàn)相關(guān)的算子流程。Pytorch原生的算子一般按照功能實(shí)現(xiàn)在一起。例如激活函數(shù)都放在Activation.h與Activation.cpp中。所以leaky_relu的實(shí)現(xiàn)就在aten/src/Aten/native/目錄下的Activation.h與Activation.cpp文件中。如下代碼段所示。不過這里定義的實(shí)現(xiàn)只是一個(gè)封裝，沒有真正的實(shí)現(xiàn)。leak_relu調(diào)用了leaky_relu_stub方法，leak_relu_backward調(diào)用了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文件內(nèi)容 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兩個(gè)函數(shù)的實(shí)現(xiàn)流程都在aten/src/Aten/native/cpu/Activation.cpp中。并且增加了兩個(gè)DISPATH（函數(shù)分發(fā)的說明）。如下代碼段所示：

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兩個(gè)函數(shù)的實(shí)現(xiàn)流程都在aten/src/Aten/native/cuda/Activation.cu中。并且增加了兩個(gè)DISPATH（函數(shù)分發(fā)的說明）。如下代碼段所示：

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

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

?

下面以實(shí)現(xiàn)一個(gè)自定義的xxx算子為例，為了簡(jiǎn)單起見，只實(shí)現(xiàn)該算子的CPU前向算子。首先在native_functions.yaml文件中增加xxx算子的描述：

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

然后在同級(jí)目錄下實(shí)現(xiàn)算子的表層實(shí)現(xiàn)文件，同樣為了簡(jiǎn)單起見，直接實(shí)現(xiàn)在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中實(shí)現(xiàn)真正的xxx_stub方法。為了簡(jiǎn)單，不做任何數(shù)值操作，只是調(diào)用printf打印相關(guān)信息。

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

編譯之后，對(duì)xxx算子進(jìn)行測(cè)試，如下所示：

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

3、C++?extention方式

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

算子名稱為lltm，首先看一下目錄結(jié)構(gòu)：

在lltm.cpp中編寫前向和反向函數(shù)的功能實(shí)現(xiàn)：

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

另外需要增加pybind11的綁定說明。因?yàn)閜ytorch的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，用于編譯生成相關(guān)的模塊

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

完成上述代碼的編寫之后，執(zhí)行：python?setup.py?install即可完成編譯。生成pytorch中可用的lltm算子。下面對(duì)新增加的lltm算子進(jìn)行測(cè)試，發(fā)現(xiàn)pytorch已經(jīng)可以準(zhǔn)確識(shí)別該算子了。

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方式能夠比較方便的增加底層算子。但是也存在一點(diǎn)缺陷。首先它是作為一個(gè)額外的擴(kuò)展模塊接入pytorch，所以在調(diào)用這些方法的時(shí)候，都是需要直接導(dǎo)入方法名稱。即無法通過torch.xxx或者tensor.xxx的方式進(jìn)行調(diào)用，另外只能支持現(xiàn)有平臺(tái)，無法擴(kuò)展到新的硬件平臺(tái)。所以Pytorch還提供了一種更加強(qiáng)大的算子擴(kuò)展能力，就是OP Register（算子注冊(cè)）方式。同樣，該方式與pytorch源碼解耦，增加和修改算子不需要重新編譯pytorch源碼。關(guān)于該部分的說明，pytroch的官方文檔中并沒有找到相關(guān)信息，但是在pytroch源碼的aten/src/ATen/core/op_registration/README.md中有一些介紹。（備注：雖然該方法與pytorch本身解耦，如果需要增加新硬件平臺(tái)對(duì)應(yīng)的算子，那么需要首先在pytroch源碼中增加對(duì)新硬件的支持，以及算子分發(fā)的DISPATH_KEY等相關(guān)信息，然后才能使用該方法注冊(cè)基于該新硬件的算子）

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

//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()));

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

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

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。