算子代码实现
简介
AI CPU算子的实现包括两部分:
- 头文件(.h文件):进行算子类的声明,自定义算子类需要继承CpuKernel基类。
- 源文件(.cc文件):重写算子类中的Compute函数,进行算子计算逻辑的实现。
头文件代码模块介绍
用户需要在算子工程的“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
- 引入相关头文件。
- 头文件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类的属性定义以及方法。
- 进行算子类的声明,此类为CpuKernel类的派生类,并需要声明重载函数Compute,Compute函数需要在算子实现文件中进行实现,详细请参见源文件代码模块介绍。算子类的声明需要在命名空间“aicpu”中,命名空间的名字“aicpu”为固定值,不允许修改。
源文件代码模块介绍
用户需要在算子工程的“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
- 引入相关头文件。
头文件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类的属性定义以及方法。
- 定义命名空间,声明常量字符指针指向算子的OpType。
namespace { const char *RESHAPE_CUST = "ReshapeCust"; }
其中,ReshapeCust为算子的OpType,RESHAPE_CUST为声明的指向算子OpType的常量指针。
- 定义命名空间aicpu,并在命名空间aicpu中实现自定义算子的Compute函数,定义算子的计算逻辑。
命名空间的名称aicpu为固定值,基类及相关定义都在aicpu命名空间中。
- Compute函数声明。
uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)
ReshapeCustCpuKernel为头文件中定义的自定义算子类,形参CpuKernelContext为CPU Kernel的上下文,包括算子的输入输出Tensor以及属性等相关信息。
- Compute函数体中,根据算子开发需求,编写相关代码实现获取输入tensor相关信息,并根据输入信息组织计算逻辑,得出输出结果,并将输出结果设置到输出tensor中。
- Compute函数声明。
- 注册算子的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算子开发