• 问答
  • 技术
  • 实践
  • 资源
推理实践落地 | 最详细的 Pytorch 底层算子扩展总结(文末附源码)
技术讨论
来源:DL工程实践


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实现,然后添加配置信息,就可以自动生成:[http://torch.xxx()](https://link.zhihu.com/?target=http://torch.xxx())、[http://torch.nn.functional.xxx()](https://link.zhihu.com/?target=http://torch.nn.functional.xxx())以及[http://tensor.xxx()](https://link.zhihu.com/?target=http://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) -> Tensor
  use_c10_dispatcher: full
  python_module: nn
  dispatch:
    CPU: leaky_relu
    CUDA: leaky_relu
    QuantizedCPU: quantized_leaky_relu

- func: leaky_relu_backward(Tensor grad_output, Tensor self, Scalar negative_slope, bool self_is_result) -> Tensor
  use_c10_dispatcher: full
  python_module: nn

其次,需要在配置文件tools/autograd/derivatives.yaml中添加算子和反向算子的对应关系,如下代码段表示,即说明了leaky_relu的反向传播函数为leaky_relu_backward。

- name: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor
  self: 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) -> Tensor
  use_c10_dispatcher: full
  python_module: nn
  dispatch:
    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,所以在调用这些方法的时候,都是需要直接导入方法名称。即无法通过http://torch.xxx或者http://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 底层代码?这份内核机制简介送给你
PyTorch 扩展自定义 PyThon/C++(CUDA) 算子的若干方法总结

  • 1
  • 0
  • 2551
收藏
暂无评论