算子代码实现

简介

AI CPU算子的实现包括两部分：

头文件代码模块介绍

用户需要在算子工程的“cpukernel/impl/reshape_cust_kernel.h”文件中进行算子类的声明，代码模块介绍如下所示：

#ifndef _AICPU_RESHAPE_CUST_KERNELS_H_
#define _AICPU_RESHAPE_CUST_KERNELS_H_

#include "cpu_kernel.h"        // CpuKernel基类以及注册宏定义
// 以下头文件根据算子需求引入 
#include "cpu_tensor.h"        // Tensor定义以及相关方法
#include "cpu_tensor_shape.h"  // Tensor shape的定义以及相关方法
#include "cpu_types.h"         // 数据类型以及格式等定义
#include "cpu_attr_value.h"    // AttrValue定义及相关方法

namespace aicpu {        // 定义命名空间aicpu
class ReshapeCustCpuKernel : public CpuKernel {        // ReshapeCust算子类继承CpuKernel基类
public:
    ~ReshapeCustCpuKernel() = default;
    virtual uint32_t Compute(CpuKernelContext &ctx) override;    // 声明函数Compute，且Compute函数需要重写
};
} // namespace aicpu
#endif

源文件代码模块介绍

用户需要在算子工程的“cpukernel/impl/reshape_cust_kernel.cc”文件中进行算子的计算逻辑实现，代码模块介绍如下所示：
#include "reshape_cust_kernel.h"     // 引入声明ReshapeCust算子类的头文件

namespace {          
const char *RESHAPE_CUST = "ReshapeCust";      //算子的OpType为ReshapeCust
}

namespace aicpu {    // 定义命名空间aicpu
uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)  // 实现自定义算子类的Compute函数
{
    ...
    return 0;
}

REGISTER_CPU_KERNEL(RESHAPE_CUST, ReshapeCustCpuKernel);  // 注册ReshapeCust算子实现
} // namespace aicpu
  1. 引入相关头文件。

    头文件reshape_cust_kernel.h，头文件代码模块介绍中声明的头文件。

    以下头文件，若在头文件代码模块介绍已经引入，无需在重复引入。

    • 头文件cpu_kernel.h，包含AI CPU算子基类CpuKernel的定义，以及Kernels的注册宏的定义。
    • 头文件cpu_tensor.h，包含AI CPU的Tensor类的定义及相关方法。
    • 头文件cpu_tensor_shape.h，包含AI CPU的TensorShape类及相关方法。
    • 头文件cpu_types.h，包含AI CPU的数据类型以及格式等定义。
    • 头文件cpu_attr_value.h，包含AttrValue类的属性定义以及方法。
  2. 定义命名空间，声明常量字符指针指向算子的OpType。

    如下所示：

    namespace {          
const char *RESHAPE_CUST = "ReshapeCust";
}

    其中，ReshapeCust为算子的OpType，RESHAPE_CUST为声明的指向算子OpType的常量指针。

  3. 定义命名空间aicpu，并在命名空间aicpu中实现自定义算子的Compute函数，定义算子的计算逻辑。
    命名空间的名称aicpu为固定值，基类及相关定义都在aicpu命名空间中。
    • Compute函数声明。
      uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)

      ReshapeCustCpuKernel为头文件中定义的自定义算子类，形参CpuKernelContext为CPU Kernel的上下文，包括算子的输入输出Tensor以及属性等相关信息。

    • Compute函数体中，根据算子开发需求，编写相关代码实现获取输入tensor相关信息，并根据输入信息组织计算逻辑，得出输出结果，并将输出结果设置到输出tensor中。
  4. 注册算子的Kernel实现。
    REGISTER_CPU_KERNEL(RESHAPE_CUST, ReshapeCustCpuKernel);
    • 第一个参数RESHAPE_CUST为2中定义的指向算子OpType的字符串指针。
    • 第二个参数ReshapeCustCpuKernel为自定义算子类的名称。

Compute函数实现

AI CPU算子实现的关键代码是Compute的实现，ReshapeCust算子的功能是将输入tensor的数据拷贝到输出tensor中，输出的tensor的shape信息将在算子原型定义的Infershape中进行推导。

算子实现代码示例如下所示：
namespace aicpu {
uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)
{
    Tensor *inputTensor = ctx.Input(0);    // 可根据获取到的输入input_tensor获取输入的shape,数据等信息
    if (inputTensor == nullptr) {
        return -1;
    }

    Tensor *outputTensor = ctx.Output(0);    // 可根据获取到的输出output_tensor获取输出的shape,数据等信息
    if (outputTensor == nullptr) {
        return -1;
    }

    // 获取输入tensor的数据地址
    auto inputData = inputTensor->GetData();
    if (inputData == nullptr) {
        return -1;
    }

    // 获取输出tensor的数据地址
    auto outputData = outputTensor->GetData();
    if (outputData == nullptr) {
        return -1;
    }

    // 将输入tensor的数据拷贝至输出tensor中
    uint64_t inputDataSize = inputTensor->GetDataSize();
    memcpy(outputData, inputData, inputDataSize);
    return 0;
}
父主题： AI CPU算子开发