pytorch自定义算子 native_functions.yaml

pytorch自定义算子 native_functions.yaml

在pytorch的文件夹中搜索native_functions.yaml，可以看到，所有pytorch原生的函数都是在这里注册的

在同文件夹下，存在一个README.md，这是官方教程，教我们怎么用native_functions.yaml注册自己自定义的算子

下面我仿照pytorch原生的relu算子，写一个自己的myrelu

首先，找到native_functions.yaml，有几个函数是我们可以参考的

- func: relu(Tensor self) -> Tensor
  use_c10_dispatcher: full
  variants: function, method
  dispatch:
    CPU, CUDA: relu
    MkldnnCPU: mkldnn_relu
    QuantizedCPU: relu_quantized_cpu

这里除了dispatch其他可以照搬。dispatch主要是针对不同的backend调用不同的kernel，比如GPU上的调用CUDA的kernel。这里展示的是CPU的，在native_functions.yaml加上

- func: myrelu(Tensor self) -> Tensor
  use_c10_dispatcher: full
  variants: function, method
  dispatch:
    CPU: myrelu

接下来，找到原生relu的实现，在pytorch/aten/src/ATen/native/Activation.cpp里面
relu的实现调用了同文件下的threshold，threshold的功能是实现 x > thresh ? x : result
模仿写这个调用太复杂了，而且如果在c++段调用函数的话，真正执行的时候每次调用函数都会经过dispatch，会浪费挺多时间的。因此，我参考另外一个激活函数写自己的myrelu

Tensor hardsigmoid(const Tensor& self) {
   
  Tensor result;
  auto iter = TensorIterator::unary_op(result, self);
  hardsigmoid_stub(iter.device_type(), iter);
  return iter.output();
}

hardsigmoid是用直线去近似sigmoid函数，它在这个文件中的实现和我们对relu的实现需求很像。这里大致做的事情是，创建一个空的返回对象，然后创建一个TensorIterator，对象类型是单操作数，然后调用hardsigmoid_stub进行具体的实现，最后返回。
在刚刚的activation.cpp中加入

Tensor myrelu(const Tensor& self) {
   
  Tensor result;
  auto iter = TensorIterator::unary_op(result, self);
  myrelu_stub(iter.device_type(), iter);
  return iter.output();
}

接下来，我们要知道还需要加上什么使得myrelu_stub也能像hardsigmoid_stub一样可以用。搜索hardsigmoid_stub，发现文件中还存在一个地方

DEFINE_DISPATCH(hardsigmoid_stub);

这段根据字面意思是说，这是一个分发的对象。
加上

DEFINE_DISPATCH(myrelu_stub);

理论上，activation.h也要看一下，activation.cpp和.h是一个整体
在h文件中存在

using hardsigmoid_fn = void(*)(TensorIterator&);
DECLARE_DISPATCH(hardsigmoid_fn, hardsigmoid_stub);

仿照，在对应的地方分别加上

using myrelu_fn = void(*)(TensorIterator&);
DECLARE_DISPACTH(myrelu_fn, myrelu_stub);

这样感觉算是告一段落，那么问题来了，myrelu_stub应该在哪里实现呢
在刚才activation.cpp的文件夹下，还有一个叫做cpu的文件夹，里面存在