aot类型自定义算子进阶用法

概述

aot类型的自定义算子采用预编译的方式，要求网络开发者基于特定接口，手写算子实现函数对应的源码文件，并提前将源码文件编译为动态链接库，然后在网络运行时框架会自动调用执行动态链接库中的函数。aot类型的自定义算子支持GPU平台的CUDA语言，和CPU平台的C和C++语言。关于aot类型的自定义算子开发的基础知识请参考基础教程。

本教程中，我们将展示aot类型自定义算子的进阶功能，包括

• aot类型自定义算子的自编译功能；

• aot类型自定义算子的属性和中间变量；

• aot类型自定义算子的动态shape支持。

对于下面用例的完整代码，请查阅这里。

aot类型自定义算子进阶用法特性简介

aot类型自定义算子的自动编译

当用户的aot类型自定义算子文件为单一文件，且编译时不需要自定义的编译选项时，可以使用自动编译功能。如此，用户可以给自定义算子提供算子实现的源文件，MindSpore会自动把源文件编译成二进制库进行调用。当前该功能支持基于GCC的C++文件编译和基于NVCC的CUDA文件编译。在使用自动编译功能的时候，有如下几点需要说明：

• MindSpore识别自动编译的方式为文件名后缀。为了使用自动编译功能，请使用后缀为cpp, cc或者cu的源文件。其他情况MindSpore将处理为二进制库的路径。

• 自动编译的结果在文件夹akg_kernel_meta下。

• 默认编译选项为：

  • C++: g++ -std=c++17 --shared -fPIC -D_GLIBCXX_USE_CXX11_ABI=0 -I./ -o $object_path, $source_path

  • CUDA 10: nvcc --shared -Xcompiler -fPIC -O3 -gencode arch=compute_70, code=sm_70 --use_fast_math --expt-relaxed-constexpr -D_GLIBCXX_USE_CXX11_ABI=0 -I./ -o $object_path, $source_path

  • CUDA 11（或者更高版本）: nvcc --shared -Xcompiler -fPIC -O3 -gencode arch=compute_80, code=sm_80 --use_fast_math --expt-relaxed-constexpr -D_GLIBCXX_USE_CXX11_ABI=0 -I./ -o $object_path, $source_path

• 由于MindSpore需要使用-D_GLIBCXX_USE_CXX11_ABI=0的编译选项，GPU平台下请避免使用版本低于10.1.168的CUDA软件栈。

aot类型自定义算子的属性和中间变量

常用的算子当中，不少算子带有属性，比如convlution的kernel size、padding和strides。带有不同属性值的算子有着相同的计算逻辑，唯一的区别是初始化时赋予属性不同的数值。此外，在算子的计算过程中，可能需要一些额外的内存空间储存中间变量。下面的计算为例，如果我们考虑input_1和input_2计算output如下公式：

tmp = Add(input_1, input_2)
output = ReduceSum(tmp, axis, keep_dims)

这里，我们需要在算子中添加如下中间变量和属性以在计算函数中使用，包括

• tmp为中间变量，记录加法的中间结果；

• axis是类型为int的属性，keep_dims是类型为bool的属性。

aot类型的自定义算子提供属性功能，如此，我们可以通过一套源码定义一类自定义算子。这类有着相同的计算逻辑，而通过算子初始化的时候对属性赋值达到不同的计算效果。此外，为了让MindSpore统一管理内存的分配和释放，aot类型的自定义算子提供了接口，指定中间变量占内存的大小，由MindSpore申请内存供计算使用。

aot类型自定义算子的动态shape支持

动态Shape，指的是算子输入或者输出的形状依赖于具体的运算，无法在编译期提前计算得出。具体来说分两种情况：算子输入的形状在编译期未知和算子输出的形状依赖具体输入的值。算子输入的形状在编译期未知的场景较为常见。任何算子，无论其计算逻辑如何，只要在支持动态shape输入的网络中使用，都需要支持这种场景。

当前自定义算子aot模式支持算子输入的形状在编译期未知的动态shape场景，通过定义c++版本的shape推导函数支持自定义算子该场景下的类型推导。

值得注意的是，目前自定义算子尚不支持算子输出的形状依赖具体输入的值的动态shape场景。

aot类型自定义算子进阶用法接口简介

主函数

源码文件中，算子实现函数的主函数必须满足如下规范：

extern "C" int FuncName(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);

其中，函数名FuncName可替换成任意有效函数名。返回值为int类型，约定0表示正常退出，非0表示发生异常。参数列表的含义如下：

• nparam (int): 输入，输出和中间变量总数。比如算子有2个输入，1个输出，1个中间变量，则nparam的值为4。

• params (void **): 输入，输出和中间变量指针数组。比如算子有2个输入，1个输出，1个中间变量，那么params[0]指向第一个输入数据，params[1]指向第二个输入数据的内存，params[2]指向输出数据的内存，params[3]指向中间变量的内存。

• ndims (int *): 输入，输出和中间变量shape维度数组。比如params[i]是个shape[1024, 1024]的张量，则ndims[i]的值为2。

• shapes (int64_t **): 输入，输出和中间变量shape数组。比如params[i]是个shape[1024, 1024]的张量，则shapes[i][0]的值为1024，shapes[i][1]的值为1024。

• dtypes (const char **): 输入，输出和中间变量数据类型数组。dtypes里的元素取值可为：”float32”、”float16”、”float”、”float64”、”int”、”int8”、”int16”、”int32”、”int64”、”uint”、”uint8”、”uint16”、”uint32”、”uint64”和”bool”。

• stream (void *): CUDA流指针，仅定义GPU算子实现时需要。

• extra_void (void *): 属性相关数据结构指针。

初始化函数

为了支持算子属性和中间变量，我们需要定义算子初始化函数。算子初始化函数定义必须满足如下规范：

extern "C" int FuncNameInit(int *ndims, int64_t **shapes, const char **dtypes, AotExtra *extra);

其中，函数名FuncName为算子主函数的名字。返回值为int类型，约定0表示正常退出，非0表示发生异常。参数列表的含义如下：

• ndims (int *): 输入输出shape维度数组。

• shapes (int64_t **): 输入输出shape数组。

• dtypes (const char **): 输入输出数据类型数组。

• extra (AotExtra *): 用于带属性的自定义算子扩展。其中AotExtra类型定义在MindSpore提供的头文件custom_aot_extra.h。

Shape推导函数

为了支持动态shape，aot类型的自定义算子中需要加入C++版本的shape推导函数。算子shape推导函数定义必须满足如下规范：

extern "C" std::vector<int64_t> FuncNameInferShape(int *ndims, int64_t **shapes, AotExtra *extra)

其中，函数名FuncName为算子主函数的名字。返回值为std::vector<int64_t>类型，为输出的shape。参数列表的含义如下：

• ndims (int *): 输入shape维度数组。

• shapes (int64_t **): 输入shape数组。

• extra (AotExtra *): 用于带属性的自定义算子扩展。其中AotExtra类型定义在MindSpore提供的头文件custom_aot_extra.h。

算子属性注册（Python）

算子属性的在初始化时的赋值通过算子注册文件实现。对于每一个属性，我们为算子注册文件创建一个attr，设置属性名和属性的值。其注册方法为

def attr(self, name=None, param_type=None, value_type=None, default_value=None, **kwargs)

其参数含义参见CustomRegOp相关接口文档。其中，在aot类型自定义算子注册时，我们注册时需要注意一下四个参数：

• name: aot类型自定义算子的属性的名称；

• param_type: 属性的参数类型。对于aot类型自定义算子的属性，这个输入固定为”required“，即必选参数；

• value_type: 属性的数值类型。对于aot类型自定义算子的属性，这个输入可以为具体的数值类型，也可以是”all”，即不限定类型；

• 最后一个输入需要指定输入名为value=，输入的值为属性的值。

aot类型自定义算子进阶用法用例

下面我们用一个Add和ReduceSum的融合算子用例来介绍aot类型自定义算子的进阶用法。该算子先把两个输入相加，在对某个轴计算求和操作，其基本计算逻辑如下：

tmp = Add(input_1, input_2)
output = ReduceSum(tmp, axis, keep_dims)

这里，我们需要在算子中添加如下中间变量和属性以在计算函数中使用，包括

• tmp为中间变量，记录加法的中间结果；

• axis是类型为int的属性，keep_dims是类型为bool的属性。

算子实现文件（C++/CUDA）:kernel.cc

为了实现算子，我们创建源文件kernel.cc，包括以下一个算子属性类add_reduce_kernel_attr和三个函数：CustomKernelInit、CustomKernelInferShape和CustomKernel。

算子属性类

首先我们定义一个数据结构贮存算子属性，该数据接口继承自AotKernelData。AotKernelData是自定义算子属性数据结构的统一基类，通过下载MindSpore提供的头文件custom_aot_extra.h放在源文件同一目录下并在文件前#include "custom_aot_extra.h"便可以使用相关接口。

#include <vector>
#include "custom_aot_extra.h"
class add_reduce_kernel_attr : public AotKernelData {
 public:
  int64_t axis;
  bool keep_dim;
};

这里我们在属性类add_kernel定义了：

• axis：成员变量，类型为int64_t；

• keep_dim：成员变量，类型为bool；

算子初始化函数

定义完算子属性类后，我们定义算子初始化函数。值得注意是，这里的初始化函数名CustomKernelInit对应，那么下面对应函数的前缀应该都为CustomKernel。

extern "C" int CustomKernelInit(int *ndims, int64_t **shapes, const char **dtypes, AotExtra *extra) {
  size_t workspace_size = 1;
  for (size_t i = 0; i < ndims[0]; i++) {
    workspace_size *= shapes[0][i];
  }

  std::vector<size_t> workspace = {workspace_size * sizeof(float)};
  extra->SetWorkSpace(workspace);

  add_reduce_kernel_attr *kernel_data_ptr = new add_reduce_kernel_attr;
  kernel_data_ptr->axis = extra->Attr<int64_t>("axis");
  kernel_data_ptr->keep_dim = extra->Attr<bool>("keep_dim");
  extra->SetKernelData(kernel_data_ptr);
  return 0;
}

这里我们需要一个中间变量workspace记录加法的中间结果，操作方式如下：

1. 计算workspace需要的内存大小：这里workspace的shape和第一个输入一样，因此先用workspace_size *= shapes[0][i]计算出workspace中元素的个数，再用workspace_size * sizeof(float)计算内存大小（这里默认元素类型为float);

2. 把所有中间变量的内存大小储存在一个std::vector<size_t>类型的对象内：std::vector<size_t> workspace = {workspace_size * sizeof(float)};。这里因为只有一个中间变量，该向量只有一个元素；

3. 通过AotExtra *extra的SetWorkSpace设置中间变量内存大小：extra->SetWorkSpace(workspace)。

另外我们需要获得两个属性axis和keep_dim的值，操作方式如下：

1. 创建一个add_reduce_kernel_attr对象指针：add_reduce_kernel_attr *kernel_ptr = new add_reduce_kernel_attr。

2. 从extra中获取对应属性的值贮存在kernel_ptr中的成员变量中：kernel_data_ptr->axis = extra->Attr<int64_t>("axis"); kernel_data_ptr->keep_dim = extra->Attr<bool>("keep_dim");。这里reduce_axis和keep_dim分别为int和bool类型，我们用extra->Attr<T>(std::string name)接口的对应模板获取该类型属性的值。

   • 这里T支持类型为：bool、string、int64_t、float、std::vector<int64_t>、std::vector<float>、std::vector<std::vector<int64_t>>和std::vector<std::vector<float>>。

3. 把kernel_ptr存在extra中供算子计算时使用：extra->SetKernelData(kernel_ptr)。

算子Shape推导函数

为了定义动态shape场景，我们定义C++版本的算子Shape推导函数如下。值得注意是，这里的算子Shape推导函数名CustomKernelInferShape和上面的初始化函数名CustomKernelInit的前缀均为前缀CustomKernel。

#include <vector>
#include "custom_aot_extra.h"

extern "C" std::vector<int64_t> CustomKernelInferShape(int *ndims, int64_t **shapes, AotExtra *extra) {
  const int64_t kDynRankSize = -2;

  if (shapes[0][0] == kDynRankSize) {
    return std::vector<int64_t>{shapes[0][0]};
  }
  int64_t axis = extra->Attr<int64_t>("axis");
  bool keep_dim = extra->Attr<bool>("keep_dim");
  if (keep_dim) {
    if (axis == 0) {
      return std::vector<int64_t>{1, shapes[0][1]};
    } else {
      return std::vector<int64_t>{shapes[0][0], 1};
    }
  } else {
    return std::vector<int64_t>{shapes[0][1 - axis]};
  }
}

在上面的例子中，我们要注意：

• 根据MindSpore的规范，动态shape输入分为dynamic shape和dynamic rank两种情况，对应的shape输入分别为：

  • dynamic shape：输入的某一维的大小未知，用-1表示。例如输入的shape为[1024, -1, 1024]，表示输入为一个三维张量，第一维和第三维长度为1024，第二维长度位置；

  • dynamic rank：输入的维度的个数位置，输入的shape固定为[-2, ]。

• 为了支持C++的shape推导函数，需要处理输入为dynamic shape和dynamic rank的场景。例如上面的例子，如果输入为dynamic rank，那么输出也是dynamic rank。因此我们判断输入为[-2, ]时，直接返回[-2, ]。

• 对于输出shape依赖属性的场景，可以通过extra->Attr<T>(std::string name)模板接口获取属性。

算子计算函数（主函数）

算子计算函数的接口规范和不带属性的自定义算子一样。值得注意是，这里的算子主函数名CustomKernel需要和上面的初始化函数名CustomKernelInit及算子Shape推导函数名CustomKernelInferShape对应。主函数和上面两个函数一起组成源文件kernel.cc。

extern "C" int CustomKernel(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream,
                         void *extra_void) {
  constexpr int OUTPUT_INDEX = 2;

  float *input_1 = static_cast<float *>(params[0]);
  float *input_2 = static_cast<float *>(params[1]);
  float *output = static_cast<float *>(params[2]);
  float *tmp = static_cast<float *>(params[3]);

  // Add
  int in_size = 1;
  for (int i = 0; i < ndims[OUTPUT_INDEX]; i++) {
    in_size *= shapes[OUTPUT_INDEX][i];
  }

  for (int i = 0; i < in_size; i++) {
    tmp[i] = input_1[i] + input_2[i];
  }

  // ReduceSum
  AotExtra *extra = static_cast<AotExtra *>(extra_void);
  auto kernel_ptr = static_cast<add_reduce_kernel_attr *>(extra->KernelData());
  bool keep_dim = kernel_ptr->keep_dim;
  int64_t axis = kernel_ptr->axis;
  int64_t input_dim_1 = shapes[0][1];
  int size;
  if (keep_dim) {
    size = shapes[1][0] * shapes[1][1];
  } else {
    size = shapes[1][0];
  }

  int ext = shapes[0][axis];
  for (int i = 0; i < size; i++) {
    output[i] = 0;
    for (int j = 0; j < ext; j++) {
      int idx = input_dim_1 * (i * axis + j * (1 - axis)) + i * (1 - axis) + j * axis;
      output[i] = output[i] + tmp[idx];
    }
  }
  return 0;
}

在计算Add时我们使用了算子的中间变量，操作如下：

1. 把params数组中的指针依次类型转化为float *。根据上面接口的介绍，数组中的元素依次为：两个输入的地址指针(input_1和input_2)，一个输出的地址指针(output)，以及一个中间变量的地址指针(tmp)；

2. 把两个输入相加的结果存在中间变量中：tmp[i] = input_1[i] + input_2[i]。

在计算ReduceSum时我们使用了算子的属性值，操作如下：

1. 把extra_void类型转化为AotExtra类型指针：AotExtra *extra = static_cast<AotExtra *>(extra_void)。

2. 从extra中获取初始化函数中创立的kernel_ptr对象指针：auto kernel_ptr = static_cast<add_reduce_kernel_attr *>(extra->KernelData())。这里extra->KernelData()获得的是一个void对象指针，需要再进一步类型转化为kernel_ptr对象指针。

3. 使用kernel_ptr中储存的属性值进行计算：bool keep_dim = kernel_ptr->keep_dim; int64_t axis = kernel_ptr->axis;。这里我们从kernel_ptr获得变量keep_dim和axis进行计算。

算子定义文件:test_custom_aot.py

为了在MindSpore中添加aot类型的自定义算子调用上面函数，我们创建文件test_custom_aot.py。

import numpy as np
from mindspore import context, Tensor
from mindspore.common import dtype as mstype
from mindspore.nn import Cell
import mindspore.ops as ops
from mindspore.ops import DataType, CustomRegOp

class ReduceDynNet(Cell):
    def __init__(self, out_types, axis, keep_dim):
        super(ReduceDynNet, self).__init__()
        reduce_cpu_info = CustomRegOp("reduce_kernel_cpu") \
            .input(0, "x1") \
            .input(0, "x2") \
            .output(0, "y") \
            .dtype_format(DataType.None_None, DataType.None_None, DataType.None_None) \
            .attr("axis", "required", "all", value=axis) \
            .attr("keep_dim", "required", "all", value=keep_dim) \
            .target("CPU") \
            .get_op_info()
        # 由于上面定义了C++版本的shape推导函数，这里的ouptut_shape可以为`None`
        self.program = ops.Custom("./kernel.cc:CustomKernel", None, out_types, "aot", reg_info=reduce_cpu_info)

    def construct(self, x, y):
        return self.program(x, y)

该文件中的ReduceDynNet包括算子注册和算子定义两个部分。

算子注册

算子属性的在初始化时的赋值通过算子注册文件实现。关于自定义算子注册的函数，参见CustomRegOp相关文档。对于每一个属性，我们为算子注册文件reduce_cpu_info创建一个attr，设置属性名和属性的值。

这里每一个attr项有四个输入：第一个为名字，如"axis"或"keep_dim"；中间两个为"required"和"all"；最后一个输入需要指定输入名为value=，输入的值为属性的值，例如这里value=axis和value=keep_dim。这里我们从网络的输入确定这两个参数的值，这两个值应该和上面初始化函数和shape推导函数中使用的extra->Attr<T>模板接口的类型匹配。

此外，如果我们需要定义多个算子注册文件，需要使用不同的算子文件名，即CustomRegOp的入参，这里为"add_with_attr_kernel_cpu"。如果需要定义另一个算子原型相同但是属性值不同的算子时，该名字不能重复。

算子定义

上面Python文件中通过自定义算子统一接口Custom定义了aot类型的自定义算子：self.program = ops.Custom("./kernel.cc:CustomKernel", None, out_types, "aot", reg_info=reduce_cpu_info)。因为我们前面定了C++版本的shape推导函数之后，这里的ouptut_shape可以为None.

值得注意的是，这里的算子定义中我们直接使用源文件名./kernel.cc，如此我们采用MindSpore提供的自动编译功能。注意这个时候要保证环境中存在对应的编译器（这里为g++，gpu环境的cu文件则需要nvcc）。

算子调用

作为测试，我们给test_custom_aot.py文件添加__main__函数如下：

if __name__ == "__main__":
    shape = (4, 5)
    axis = 1
    keep_dim = False
    context.set_context(device_target="CPU")

    input_x = np.ones(shape).astype(np.float32)
    input_y = np.ones(shape).astype(np.float32)

    test = ReduceDynNet(mstype.float32, axis, keep_dim)
    dyn_x = Tensor(shape=[4, None], dtype=mstype.float32)
    # set the net to dynamic shape
    test.set_inputs(dyn_x, dyn_x)
    output = test(Tensor(input_x),Tensor(input_y))
    print(output)

执行文件调用算子：

python test_custom_aot.py

执行结果：

[10. 10. 10. 10.]