实现流程
用户需要在算子工程的“cpukernel/impl/xx.cc”文件中进行算子的计算逻辑实现,如下所示:
// 包含声明算子类的头文件 #include "sample_kernels.h" // 包含AI CPU基础库头文件 #include "cust_cpu_utils.h" namespace { //sample为算子原型中注册的算子的类型,可查看算子原型定义 const char *SAMPLE = "sample"; } // 定义命名空间aicpu namespace aicpu { // 实现自定义算子类的Compute函数 uint32_t SampleCpuKernel::Compute(CpuKernelContext &ctx) { Tensor *x0 = ctx.Input(0); // 对输入tensor进行基本校验,比如判断是否为空等操作 // 可根据获取到的输入tensor x0获取输入的shape,数据等信息 ... Tensor *x1 = ctx.Input(1); // 对输入tensor进行基本校验,比如判断是否为空等操作 // 可根据获取到的输入tensor x1获取输入的shape,数据等信息 ... Tensor *y0 = ctx.Output(0); // 可根据获取到的输出tensor y0获取输出的shape,数据等信息 ... AttrValue *attr = ctx.GetAttr(attr); //获取属性信息,并对属性进行基本校验,比如判断是否为空等操作 ... // 根据输入信息组织计算逻辑,得到输出结果,并将结果设置到输出的tensor中 ... // 根据需要调用Dump日志接口,打印相关调试信息 ... // 动态shape类算子需要额外更新输出的tensor的shape等信息 ... return 0; } // 注册该算子实现 REGISTER_CPU_KERNEL(SAMPLE, SampleCpuKernel); } // namespace aicpu
- 引入相关头文件。
头文件sample_kernels.h,算子类声明中声明所在的头文件。
- 定义命名空间,声明常量字符指针指向算子的OpType。
如下所示:
namespace { //sample为算子的OpType const char *SAMPLE = "sample"; }
其中,sample为算子原型中注册的算子的类型,可查看算子原型定义,SAMPLE为声明的指向算子OpType的常量指针。
- 定义命名空间aicpu,并在命名空间aicpu中定义算子的Compute函数,用于实现算子的计算逻辑。
命名空间的名称aicpu为固定值,基类及相关定义都在aicpu命名空间中,声明如下所示:
namespace aicpu { uint32_t SampleCpuKernel::Compute(CpuKernelContext &ctx) { ... ... }
SampleCpuKernel为头文件中定义的自定义算子类,形参CpuKernelContext为CPU Kernel的上下文,包括算子的输入输出Tensor以及属性等相关信息。
- Compute函数实现。
获取输入/输出Tensor相关信息,并进行合法性校验,然后根据输入信息组织计算逻辑,得出计算结果,并将输出结果设置到输出Tensor中。Compute函数实现的具体介绍请参考Compute函数实现。
- 注册算子的Kernel实现。
REGISTER_CPU_KERNEL(SAMPLE, SampleCpuKernel);
- 第一个参数SAMPLE为2中定义的指向算子OpType的字符串指针。
- 第二个参数SampleCpuKernel为自定义算子类的名称。
父主题: 计算逻辑实现