飞桨高阶使用教程：自定义CPU算子的实现和使用

算子（Operator，简称Op）是构建神经网络的基础组件。在网络模型中，算子对应层中的计算逻辑，例如：卷积层（Convolution Layer）是一个算子；全连接层（Fully-connected Layer， FC layer）中的权值求和过程，是一个算子。学会定制化算子的C++实现可以更深入地了解神经网络运行的底层逻辑。

张量与算子

• 张量(Tensor)，可以理解为多维数组,可以具有任意多的维度,不同Tensor可以有不同的数据类型(dtype)和形状(shape)。
• 算子(Operator)也简称为OP，负责对Tensor执行各种运算处理,可以理解为一个计算函数,函数的输入和输出都为Tensor。在PaddlePaddle中定义了大量的Operator来完成常见神经网络模型的Tensor运算处理,如conv2d, pool2d, Relu等

一、底层原理

PaddlePaddle中所有Op算子都会注册在OpInfoMap中,在Python端调用Op执行运算操作时,通过TraceOp在OpInfoMap找到对应的Op并调用其计算kernel ,完成计算并返回结果。

换句话说，因为硬件的不同，相同的算子需要不同的kernel，比如你写了一个在CPU上执行的算子，那么这个算子只能在CPU上运行，要想在别的计算平台上运行，还需要实现该平台的kernel。这篇文章讲的主要是怎么实现一个能在CPU上运行的算子。

PaddlePaddle Op算子体系(动态图模式)

例如，在动态图模式执行Y=relu(X)时,框架会通过TraceOp来完成:

1. 调用relu算子的forward计算函数完成Y的计算
2. 创建backward所需的Op算子以及输入输出变量(此时不进行计算,待后续调用backward()后才会进行反向计算)

PaddlePaddle Op算子正反向计算(动态图模式)

通过上面展示的底层原理，其实不难发现，一个算子最关键的部分就是前向传播与反向计算，这两个部分是算子的核心。

C++自定义算子内部原理(动态图模式)

在动态图模式执行Y=custom_ relu(X)时:使用C++自定义算子与原生算子的执行流程相同。

但和原生算子区别是原生算子随框架一起编译；自定义算子可单独编译。但最后都会注册到OpInfoMap中。

二、C++自定义算子格式

1.基本格式

在编写运算函数之前，需要引入 PaddlePaddle 扩展头文件：

#include "paddle/extension.h"

算子运算函数有特定的函数写法要求，在编码过程中需要遵守，基本形式如下：

std::vector<paddle::Tensor> OpFucntion(const paddle::Tensor& x, ..., int attr, ...) {
  ...
}

这一部分其实就是固定格式，所有用C++编写的Paddle算子都需要使用这个格式。换句话说，这是Paddle提供的算子接口，只需要按照这个接口定义算子即可。

2.适配多种数据类型

在实际开发中，一个算子往往需要支持多种数据类型，这时就需要用到模板类。在上面接口上方定义：

template <typename data_t>

需要注意的是：模板参数名 data_t 用于适配不同的数据类型，不可更改为其他命名，否则会编译失败

然后通过 switch-case 语句实现支持多种数据类型的操作：

switch(x.type()) {
  case paddle::DataType::FLOAT32:
    ...
    break;
  case paddle::DataType::FLOAT64:
    ...
    break;
  default:
    PD_THROW(
      "function  ... is not implemented for data type `",
      paddle::ToString(x.type()), "`");
}

如果不想使用 switch-case 来实现，也可以使用官方提供的dispatch宏，如PD_DISPATCH_FLOATING_TYPES

3.维度与类型的推导

PaddlePaddle 框架同时支持动态图与静态图的执行模式，在静态图模式下，组网阶段需要完成 Tensor shape 和 dtype 的推导，从而生成正确的模型描述，用于后续Graph优化与执行。因此，除了算子的运算函数之外，还需要实现前向运算的维度和类型的推导函数。

维度推导（InferShape）和类型推导（InferDtype）的函数写法也是有要求的，格式如下：

需要注意的是，输入输出参数与forward计算函数的输入输出Tensor应该按顺序一一对应：

对于仅有一个输入Tensor和一个输出Tensor的自定义算子，如果输出Tensor和输入Tensor的shape和dtype一致,可以省略InferShape和InferDtype函数的实现，但其他场景下均需要实现这两个函数。

4.自定义算子注册

最后，需要调用 PD_BUILD_OP 系列宏，构建算子的描述信息，并关联前述算子运算函数和维度、类型推导函数。其格式如下：

PD_BUILD_OP(op_name)
    .Inputs({"X"})
    .Outputs({"Out"})
    .SetKernelFn