算子代码实现（AI CPU）

简介

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

• 头文件：进行算子类的声明，自定义算子类需要继承CpuKernel基类。
• 源文件：重写算子类中的Compute函数，进行算子计算逻辑的实现。

头文件定义

用户需要在算子工程的“cpukernel/impl/xx.h”文件中进行算子类的声明，如下所示：

// CpuKernel基类以及注册宏定义
#include "cpu_kernel.h" 
// 定义命名空间aicpu
namespace aicpu {
// 算子类继承CpuKernel基类     
class SampleCpuKernel : public CpuKernel {  
public:
~SampleCpuKernel() = default;
// 声明函数Compute，且Compute函数需要重写
uint32_t Compute(CpuKernelContext &ctx) override;    
};
} // namespace aicpu

• 引入头文件“cpu_kernel.h”。

  头文件“cpu_kernel.h”中包含了AI CPU算子基类CpuKernel的定义。

  此头文件会自动引入如下头文件：

  • cpu_tensor.h，包含了AI CPU的Tensor类的定义及相关方法。
  • cpu_tensor_shape.h，包含了AI CPU的TensorShape类及相关方法。
  • cpu_types.h，包含了AI CPU的数据类型以及格式等定义。
  • cpu_attr_value.h，包含了AttrValue类的属性定义以及方法。

Compute函数实现

用户需要在算子工程的“cpukernel/impl/xx.cc”文件中进行算子的计算逻辑实现，如下所示：

// 引入声明算子类的头文件
#include "sample_kernels.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中
...
// 动态shape类算子需要额外更新输出的tensor的shape等信息
...
return 0;
}
// 注册该算子实现
REGISTER_CPU_KERNEL(SAMPLE, SampleCpuKernel);  
} // namespace aicpu

1. 引入相关头文件。

   头文件sample_kernels.h，头文件定义中声明的头文件。

2. 定义命名空间，声明常量字符指针指向算子的OpType。

   如下所示：

   namespace {   
   //sample为算子的OpType       
   const char *SAMPLE = "sample";      
   }

   其中，sample为算子原型中注册的算子的类型，可查看算子原型定义，SAMPLE为声明的指向算子OpType的常量指针。

3. 定义命名空间aicpu，并在命名空间aicpu中定义算子的Compute函数，用户实现算子的计算逻辑。

   命名空间的名称aicpu为固定值，基类及相关定义都在aicpu命名空间中，声明如下所示：

   namespace aicpu {
   uint32_t SampleCpuKernel::Compute(CpuKernelContext &ctx) {
   ... ...
   }

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

4. Compute函数实现。

   获取输入/输出Tensor相关信息，并进行合法性校验，然后根据输入信息组织计算逻辑，得出计算结果，并将输出结果设置到输出Tensor中。

   例如：

   // 从context中获取input tensor
   Tensor *input = ctx.Input(0);  
   
   // 对输入tensor进行基本校验
   // 例如，对获取到的input进行空指针校验
   if (input == nullptr) {    
   return 1;  
   }
   
   // 获取input tensor的shape信息
   auto inputShape = input->GetTensorShape();
   for (int32_t i = 0; i < inputShape->GetDims(); ++i) {
   std::cout << "dim[" << i << "] size:" << inputShape->GetDimSize(i) << std::endl;
   }
   
   // 获取input tensor的DataType
   DataType inputType = input->GetDataType();
   
   // 获取input tensor的数据地址
   auto inputData = input->GetData();
   
   // 获取输出tensor的数据地址以及shape
   Tensor *output = ctx.Output(0);
   auto outputShape = output->GetTensorShape();
   auto outputData = output->GetData();
   // 保存输出结果
   outputData[0] = inputData[0];

   以上算子实现时使用到API接口介绍可参见AI CPU API。

   1. 合法性校验一般包含如下几类的校验：
      • 对获取到的Input进行空指针校验，此校验必选。
      • 对输入输出个数进行校验。
      • 对算子输入的内在逻辑进行校验。

        例如，对于多输入算子，多个tensor的dtype需要保持一致，此时需要校验多个输入的dtype是否一致。若多输入的内在逻辑要求已经在算子原型定义的Verify函数中进行实现，则compute函数中的此校验可不再实现。

      • 对dtype的校验。

        可根据算子实际情况来选择是否进行dtype的校验，若某算子仅支持A、B两种数据类型，其他数据类型都不支持，此时可在实现算子计算逻辑前对dtype进行校验，判断dtype是否在支持的dtype列表中。

   2. 计算逻辑的实现。

      根据算子输入支持的数据类型分别进行计算，伪代码如下：

      // 获取第i个输入的类型
      auto data_type = ctx.Input(i)->GetDataType();  
      switch (data_type) { 
          case DT_FLOAT16:
              return OpCompute<Eigen::half>(...);
          case DT_FLOAT:
              return OpCompute<float>(...);
          case DT_DOUBLE:
              return OpCompute<double>(...);
          case DT_INT8:
              return OpCompute<int8_t>(...);
          case DT_INT16:
              return OpCompute<int16_t>(...);
        
          ... ...
      
          default:      
              return PARAM_INVAILD;      
      }

      其中OpCompute函数为算子的计算过程实现函数。

      PARAM_INVAILD的定义如下所示：

      const uint32_t PARAM_INVAILD = 1;

      算子计算逻辑实现时，有以下几点注意点：

      • 由于C++自身不支持半精度浮点类型，对于半精度数据类型的计算实现，可使用第三方库Eigen来表示（建议使用3.3.9版本），具体可参考资料：Link。

        例如，对于Less算子，输入为半精度时，用Eigen进行强转。

        auto input = reinterpret_cast<Eigen::half *>(input_0->GetData());

        说明：

        第三方Eigen库还提供了比较系统的矩阵和向量等线性代数相关的运算操作，若您的算子实现涉及到相关操作，可以借助Eigen库实现。

        例如，使用Eigen库进行矩阵的定义、初始化，并求取行列式的代码示例如下所示：

        #include "Eigen/Dense"
        int m,n;
        Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic> eMatrix(m, n);
        for (int i = 0; i < m; i++) {
          for (int j = 0; j < n; j++) {
            eMatrix(i, j) = i * m + j * n;
          }
        }
        //Using eigen to calculate the Determinant
        float result = eMatrix.determinant();

      • 对于动态shape算子，无法根据算子原型定义的InferShape推导得到输出tensor的shape, 所以需要在Compute函数中完成输出shape的计算与更新。

        std::vector<int64_t> dims = {inputData[0]，inputData[1]，3，4}
        outputShape ->SetDimSizes(dims);

        Sample仓中的UniqueCust算子为动态Shape算子，您可以参考UniqueCust算子的代码实现。

      • 算子计算过程可分块并行执行，这样可以有效地使用昇腾AI处理器的硬件资源，使性能达到最优。
        

        • 分块并行为性能提升的一个手段，开发者可根据对算子的性能要求选择是否需要进行并行计算。
        • 分块并行执行的功能，仅支持输入参数间独立运算的场景，若输入参数间存在数据依赖，则无法进行分块并行计算。
        • 使用此功能时，需要在算子信息库定义中设置“opInfo.flagSupportBlockDim”为“True”，并设置“opInfo.functionName”为“RunCpuKernelWithBlock”。
        1. 首先使用GetAttr接口获取分块数目以及本次计算的分块ID，示例如下：

           uint32_t blockdim = ctx.GetAttr("block_num")->GetInt();
           uint32_t blockid = ctx.GetAttr("block_id")->GetInt();

           说明：分块数目是系统根据用户配置的BlockDim切分原则（算子信息库中配置的opInfo.blockDimByIndex）及CPU核数自动计算的。每一个分块都会分配一个“block_id”，“block_id”的取值范围为blockdim-1。

           获取“block_num”及“block_id”，用户可自行进行一些基本校验。

        2. 计算本次计算（当前分块的计算）的偏移量及数据量。

           例如，若“opInfo.blockDimByIndex”配置为-1，即按照第一个输入参数的元素个数进行BlockDim的切分，则计算偏移量及数据量的代码示例如下：

           // 获取第一个输入参数的元素个数
           int64_t total = input0->NumElements();
           int64_t startpos = 0;
           int64_t len = total;
           if (blockdim != 1) {
             // 计算每一块的最大数据量
             uint32_t per_unit = std::ceil(total / blockdim);
             // 得出本次计算的偏移量
             startpos =  blockid * per_unit;
             // 得出本次计算的数据量。
             // blockid的取值范围为：0~blockdim-1，为避免最后一块数据存在拖尾，所以当blockid为最后一块时，len取值为total - per_unit * (blockdim - 1)
             len = blockid < blockdim - 1 ? per_unit : (total - per_unit * (blockdim - 1));
           }

        3. 进行算子的计算逻辑的实现。
           以Add算子为例，代码示例如下所示：

           for (int i = startpos; i < startpos + len; i++) {
             y[i] = x0[i] + x1[i];
           } 

        完整的分块并行计算的算子样例可参见开源Ascend Sample仓，更多样例可参见自定义算子模板。

5. 注册算子的Kernel实现。

   REGISTER_CPU_KERNEL(SAMPLE, SampleCpuKernel);

   • 第一个参数SAMPLE为2中定义的指向算子OpType的字符串指针。
   • 第二个参数SampleCpuKernel为自定义算子类的名称。

样例参考

本节设计实现一个对元素求反余弦值的Acos算子。