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
#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,头文件代码模块介绍中声明的头文件。
以下头文件,若在头文件代码模块介绍已经引入,无需在重复引入。
namespace {
const char *RESHAPE_CUST = "ReshapeCust";
}
其中,ReshapeCust为算子的OpType,RESHAPE_CUST为声明的指向算子OpType的常量指针。
uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)
ReshapeCustCpuKernel为头文件中定义的自定义算子类,形参CpuKernelContext为CPU Kernel的上下文,包括算子的输入输出Tensor以及属性等相关信息。
REGISTER_CPU_KERNEL(RESHAPE_CUST, ReshapeCustCpuKernel);
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;
}