自定义算子（Ascend）

    当开发网络遇到内置算子不足以满足需求时，你可以利用MindSpore的Python API方便快捷地扩展昇腾AI处理器的自定义算子。

    添加一个自定义算子，需要完成算子原语注册、算子实现、算子信息注册三部分工作。

    其中：

    • 算子原语：定义了算子在网络中的前端接口原型，也是组成网络模型的基础单元，主要包括算子的名称、属性（可选）、输入输出名称、输出shape推理方法、输出dtype推理方法等信息。

    • 算子实现：通过TBE（Tensor Boost Engine）提供的特性语言接口，描述算子内部计算逻辑的实现。TBE提供了开发昇腾AI芯片自定义算子的能力。你可以在页面申请公测。

    • 算子信息：描述TBE算子的基本信息，如算子名称、支持的输入输出类型等。它是后端做算子选择和映射时的依据。

    本文将以自定义Square算子为例，介绍自定义算子的步骤。

    每个算子的原语是一个继承于PrimitiveWithInfer的子类，其类型名称即是算子名称。

    自定义算子原语与内置算子原语的接口定义完全一致：

    • 属性由构造函数__init__的入参定义。本用例的算子没有属性，因此__init__没有额外的入参。带属性的用例可参考MindSpore源码中的custom add3用例。

    • 输入输出的名称通过init_prim_io_names函数定义。

    • 输出Tensor的shape推理方法在infer_shape函数中定义，输出Tensor的dtype推理方法在infer_dtype函数中定义。

    以Square算子原语cus_square.py为例，给出如下示例代码。

    通常编写一个算子的实现，需要编写一个计算函数和一个入口函数。

    算子的计算函数主要用来封装算子的计算逻辑供主函数调用，其内部通过调用TBE的API接口组合实现算子的计算逻辑。

    算子的入口函数描述了编译算子的内部过程，一般分为如下几步：

    1. 准备输入的placeholder，placeholder是一个占位符，返回一个Tensor对象，表示一组输入数据。

    2. 调用计算函数，计算函数使用TBE提供的API接口描述了算子内部的计算逻辑。

    3. 调用Schedule调度模块，调度模块对算子中的数据按照调度模块的调度描述进行切分，同时指定好数据的搬运流程，确保在硬件上的执行达到最优。默认可以采用自动调度模块（auto_schedule）。

    4. 调用cce_build_code编译生成算子二进制。

    更多关于使用TBE开发算子的内容请参考，关于TBE算子的调试和性能优化请参考MindStudio文档。

    算子信息是指导后端选择算子实现的关键信息，同时也指导后端为算子插入合适的类型和格式转换。它通过TBERegOp接口定义，通过op_info_register装饰器将算子信息与算子实现入口函数绑定。当算子实现py文件被导入时，op_info_register装饰器会将算子信息注册到后端的算子信息库中。更多关于算子信息的使用方法请参考TBERegOp的成员方法的注释说明，算子信息的字段含义可以参考。

    下面以Square算子的TBE实现square_impl.py为例进行介绍。square_compute是算子实现的计算函数，通过调用te.lang.cce提供的API描述了x * x的计算逻辑。cus_square_op_info是算子信息，通过TBERegOp来定义。

    • TBERegOp("CusSquare")中算子注册名称CusSquare需要与算子名称一致。

    • fusion_type("OPAQUE")中OPAQUE表示自定义算子采取不融合策略。

    • kernel_name("CusSquareImpl")中CusSquareImpl需要与算子入口函数名称一致。

    • dtype_format用来描述算子支持的数据类型，下面示例中注册了两项，说明该算子支持两种数据类型，每一项需按照输入和输出的顺序依次描述支持的格式。第一个dtype_format说明支持的第一种数据类型是input0为F32_Default格式，output0为F32_Default格式。第二个dtype_format说明支持的第二种数据类型是input0为F16_Default格式，output0为F16_Default格式。

    1. from __future__ import absolute_import
    2. from te import tvm
    3. from topi import generic
    4. import te.lang.cce
    5. from topi.cce import util
    6. from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
    9.     """
    10.     The compute function of the CusSquare implementation.
    11.     """
    12.     res = te.lang.cce.vmul(input_x, input_x)
    13.     return res
    15. # Define the kernel info of CusSquare.
    16. cus_square_op_info = TBERegOp("CusSquare") \
    17.     .partial_flag(True) \
    18.     .async_flag(False) \
    19.     .binfile_name("square.so") \
    20.     .compute_cost(10) \
    21.     .kernel_name("CusSquareImpl") \
    22.     .input(0, "x", False, "required", "all") \
    23.     .output(0, "y", False, "required", "all") \
    24.     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
    25.     .dtype_format(DataType.F16_Default, DataType.F16_Default) \
    26.     .get_op_info()
    28. # Binding kernel info with the kernel implementation.
    29. @op_info_register(cus_square_op_info)
    30. def CusSquareImpl(input_x, output_y, kernel_name="CusSquareImpl"):
    31.     """
    32.     The entry function of the CusSquare implementation.
    33.     """
    34.     shape = input_x.get("shape")
    35.     dtype = input_x.get("dtype").lower()
    37.     shape = util.shape_refine(shape)
    38.     data = tvm.placeholder(shape, name="data", dtype=dtype.lower())
    40.     with tvm.target.cce():
    41.         res = square_compute(data, output_y)
    42.         sch = generic.auto_schedule(res)
    44.     config = {"print_ir": False,
    45.               "name": kernel_name,
    46.               "tensor_list": [data, res]}
    48.     te.lang.cce.cce_build_code(sch, config)

    自定义算子与内置算子在网络中的使用方法一样，通过导入原语直接使用。下面以CusSquare的单算子网络测试为例进行说明。

    在test_square.py文件中定义网络。

    1. import numpy as np
    2. import mindspore.nn as nn
    4. from mindspore import Tensor
    5. # Import the definition of the CusSquare primtive.
    6. from cus_square import CusSquare
    7. context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
    9.     def __init__(self):
    10.         super(Net, self).__init__()
    11.         self.square = CusSquare()
    13.     def construct(self, data):
    14.         return self.square(data)
    16. def test_net():
    17.     x = np.array([1.0, 4.0, 9.0]).astype(np.float32)
    18.     square = Net()
    19.     output = square(Tensor(x))
    20.     print("x: ", x)
    21.     print("output: ", output)

    执行用例:

    执行结果:

    1. x: [1. 4. 9.]
    2. output: [1. 16. 81.]

    如果算子要支持自动微分，需要在其原语中定义其反向传播函数（bprop）。你需要在bprop中描述利用正向输入、正向输出和输出梯度得到输入梯度的反向计算逻辑。反向计算逻辑可以使用内置算子或自定义反向算子构成。

    定义算子反向传播函数时需注意以下几点：

    • bprop函数的入参顺序约定为正向的输入、正向的输出、输出梯度。若算子为多输出算子，正向输出和输出梯度将以元组的形式提供。

    • bprop函数的返回值形式约定为输入梯度组成的元组，元组中元素的顺序与正向输入参数顺序一致。即使只有一个输入梯度，返回值也要求是元组的形式。

    例如，增加bprop后的CusSquare原语为：

    1. class CusSquare(PrimitiveWithInfer):
    2.     @prim_attr_register
    3.     def __init__(self):
    4.         """init CusSquare"""
    5.         self.init_prim_io_names(inputs=['x'], outputs=['y'])
    6.         from square_impl import CusSquareImpl
    8.     def infer_shape(self, data_shape):
    9.         return data_shape
    11.     def infer_dtype(self, data_dtype):
    12.         return data_dtype
    14.     def get_bprop(self):
    15.         def bprop(data, out, dout):
    16.             twos_like = ops.OnesLike()(data) * 2.0
    17.             gradient = ops.Mul()(data, twos_like)
    18.             dx = ops.Mul()(gradient, dout)
    19.             return (dx,)
    20.         return bprop

    在test_square.py文件中定义反向用例。

    1. pytest -s tests/st/ops/custom_ops_tbe/test_square.py::test_grad_net

    执行结果:

    1. x: [1. 4. 9.]