功能实现

导入Python模块

进行TBE DSL算子开发时，首先需要在算子实现文件中导入昇腾AI软件栈提供的Python模块，代码示例如下所示，算子实现文件的命名请参见算子定义命名规则。

如果您进行代码实现时依赖了其他自行引入的Python模块，请自行进行依赖导入。

import te.lang.cce
from te import tvm
from te.platform.fusion_manager import fusion_manager
from topi import generic

其中：

• “te.lang.cce”：引入TBE支持的特定域语言接口，包括常见的运算vmuls、vadds、matmul等。

  具体的接口定义可查看ATC安装路径下“/python/site-packages/te/lang/cce/”目录下的python函数，使用方法请参见TBE DSL API。

• “te.tvm”：引入TVM后端代码生成机制。

  具体的接口定义可查看ATC安装路径下“/python/site-packages/te/tvm”目录下的python函数，使用方法请参见https://docs.tvm.ai/。

• “te.platform.fusion_manager.fusion_manager”：提供了实现算子的UB自动融合的接口。

  具体的接口定义可查看ATC安装路径下“/python/site-packages/te/platform/fusion_manager.py”文件中的fusion_manager函数的定义。

• “topi.generic”：提供了自动算子调度接口。

  具体的接口定义可查看ATC安装路径下“/python/site-packages/topi/generic”目录下的所有python函数，使用方法请参见TBE DSL API。

算子函数声明

算子的代码实现中包括两个函数：算子接口函数与算子compute函数，算子的compute函数会在算子接口函数中被调用。

下面详细介绍这两个函数的声明规则。

• 算子接口函数声明

  如下所示，一个算子的接口函数中包含了算子输入信息、算子输出信息以及内核名称，函数的声明信息需要与算子信息定义文件中的信息对应。

  def operationname(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="KernelName")

  • 算子接口函数名称operationname当前版本请与算子实现文件名称保持一致，命名规则请参见算子定义命名规则。
  • input_x1, input_x2：算子的输入tensor，每个tensor需要采用字典的形式进行定义，包含shape、ori_shape、format、ori_format与dtype信息，例如：

    dict input_x1 = {'shape' : (2,2), 'ori_shape' : (2,2), 'format': 'ND', 'ori_format':'ND', 'dtype' : 'float16'}

    输入tensor的名称、个数及顺序需要与算子信息定义文件/tbe/op_info_cfg/ai_core/<soc_version>/operationname.ini中的定义保持一致。

  • output_y：算子的输出tensor，包含shape和dtype等信息，字典格式，此字段为预留位。
  • attribute1，attribute2...：算子的属性，此处需要为算子的属性赋默认值，算子属性的名称、个数与顺序需要与算子信息定义文件/tbe/op_info_cfg/ai_core/<soc_version>/operationname.ini中的定义保持一致。

    若算子无相关属性信息，此参数忽略。

  • kernel_name：算子在内核中的名称（即生成的二进制文件与算子描述文件的名称），用户自定义，保持唯一，只能是大小写字母、数字、“_”的组合，且必须是字母或者“_”开头，长度小于或等于200个字符。

  不带属性的sqrt算子的接口函数声明如下：

  def sqrt(input_x, output_y, kernel_name="sqrt"):

  带属性的reduce_sum算子的接口函数声明如下：

  def reduce_sum(x, y, axis=None, keep_dims=None, kernel_name="reduce_sum")

• compute函数声明

  如下所示，compute函数的入参是计算过程中涉及的所有输入tensor、attribute，和输出tensor，以及内核名称。

  @fusion_manager.register("KernelName")
  def operationname_compute(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="KernelName")

  • 装饰器@fusion_manager.register("KernelName")是算子计算实现声明中必需的，其作用是整网运行时支持算子做UB自动融合，使得当前自定义算子可以在UB中根据UB融合规则自动与其他算子的计算进行组装。
  • input_x1, input_x2：compute函数的入参，为在 算子接口函数声明中声明的输入tensor对应的placeholder，包含shape和dtype等信息。
  • output_y,attribute1=None,xxx等参数，都是从 算子接口函数声明中的算子接口函数中透传过来的，与算子接口函数的声明保持一致即可。

  例如，对于sqrt算子，compute函数定义如下：

  @fusion_manager.register("sqrt")
  def sqrt_compute(input_data, output_data, kernel_name="sqrt"):

  对于reduce_sum算子，算子接口和计算函数定义如下：

  @fusion_manager.register("reduce_sum")
  def reduce_sum_compute(x, y, axis=None, keep_dims=None, kernel_name="reduce_sum")

算子函数实现

完成算子函数声明后，就要具体实现算子接口函数和compute函数。

1. 首先在算子接口函数operatorname( )中，获取算子输入tensor的shape以及dtype，并可自行实现基本的校验功能。
   1. 获取算子输入tensor的shape以及dtype，为后续定义输入tensor的张量占位符做准备。

      def add(input_x, input_y, output_z, kernel_name="add"):
          shape_x = input_x.get("shape")
          shape_y = input_y.get("shape")
          input_data_type = input_x.get("dtype").lower()
          input_data_type_y = input_y.get("dtype").lower()

   2. （可选）在算子实现函数中添加算子输入/输出及属性基本信息校验，有助于在算子编译阶段，提前发现问题。

      例如，对于Add算子，首先校验两个输入的dtype是否一致，然后校验输入的数据类型是否在允许的数据类型列表中，代码实现如下所示。

          if input_data_type != input_data_type_y:
              raise RuntimeError(
                  "the input_x and input_y should have the same data type.")
          check_tuple = ("float16", "float32", "int32")
          if input_data_type not in check_tuple:
              raise RuntimeError("only support %s while dtype is %s" %
                                 (",".join(check_tuple), input_data_type))

      开发者可根据算子特点自定义实现校验函数。

2. 然后根据shape与dtype定义好输入tensor的张量占位符。
   例如：

   data_input = tvm.placeholder(shape, name="data_input", dtype=dtype)

   使用TVM的placeholder接口对输入tensor进行占位，返回一个tensor对象，此位置中的数据在程序运行时才被指定。

   注意：

   调度与编译中的tensor_list的输入tensor需要是tvm.placeholder接口返回的tensor对象，所以此对象在后续计算过程实现中不能被替换。

   如下所示：

       #返回占位的data_input
       data_input = tvm.placeholder(shape, name='data', dtype=dtype)      
        if dtype == "float16":
            #将data_input的类型转换为float32，然后重新赋值给data_input，此时data_input的内容已经被改变
            data_input = te.lang.cce.cast_to(data_input, "float32")        
           ......
       with tvm.target.cce():
           schedule = generic.auto_schedule(res)
       config = {"print_ir":need_print, 
                       "need_build":need_build,
                       "name":kernel_name, 
                       "tensor_list":[data_input,res]} 
       te.lang.cce.cce_build_code(schedule,config)

   以上代码中，通过data_input = te.lang.cce.cast_to(data_input, "float32")转换数据类型后，placeholder返回的data_input对象已经被覆盖，编译配置tensor_list中的data_input已经不是原placeholder接口返回的tensor，此时算子实现代码编译时会出现以下错误：

   所以可重新定义一个tensor用于存储转换数据类型后的输入进行计算，如下所示：

   data_input1 = te.lang.cce.cast_to(data_input, "float32")

   或者如步骤3所示，计算过程在compute函数中进行，将placeholder返回的输入tensor通过形参传入compute函数进行计算，会生成新的地址用于计算，也可避免placeholder返回的tensor对象被覆盖的情况。

3. 在算子接口定义函数中调用compute函数进行计算过程的描述。

   例如：

   res = add_compute(data_x, data_y, output_z, kernel_name)

   输入tensor为使用tvm.placeholder定义的占位tensor，其他为算子接口函数透传的参数。

4. 算子compute函数的实现。

   在compute函数中，完成算子的计算过程，计算过程的实现主要根据算子分析中的TBE DSL API进行代码开发。

   入门示例中介绍了简单的add算子的实现过程，下面我们以计算公式较复杂的relu算子为例，讲解算子的计算过程的实现以及部分DSL接口在使用过程中的注意事项。

   假设通过进行算子分析，得到relu算子的计算公式如下：

   计算实现代码如下所示：

   @fusion_manager.register("relu")
   def relu_compute(x, y, kernel_name="relu"):
       inp_dtype = x.dtype     # 获取输入数据的数据类型
       shape = x.shape         # 获取输入数据的形状
   
       # 若数据类型为float32、int32，使用vmax操作，避免精度损失
       if inp_dtype in ("float32", "int32"):
           tensor_zero = te.lang.cce.broadcast(tvm.const(CONST_ZERO, inp_dtype),shape)     # 返回形状与输入数据相同，每一个元素都为0，每一个元素的数据类型都为输入数据的数据类型的tensor
           data_res = te.lang.cce.vmax(x, tensor_zero)    # 取x与tensor_zero中的大值
       else:
           data_res = te.lang.cce.vrelu(x)      # 若数据类型为int8、float16，直接做relu操作。
   
       data_res = te.lang.cce.cast_to(data_res, inp_dtype)
   
       return data_res

   • 由于te.lang.cce.vrelu( )接口会将int8、uint8、int32、float32的数据类型转换为float16，而int32、float32进行数据类型转换时会造成精度损失，所以为了避免精度损失，对于这两种数据类型的输入，采用te.lang.cce.vmax( )接口取输入数据与0之间的大值。
   • TBE DSL中vmax接口要求两个输入tensor的shape相同，一般使用te.lang.cce.broadcast接口将输入tensor的shape广播到相同shape，一般取两个输入tensor的shape中每个维度的大值组成的shape。

   算子计算函数实现中的其他小技巧：

   • 若输入tensor数据类型不是float32，可以将其转换为float32进行计算，可以提高中间计算结果的精度，最后的结果输出时需要将数据类型转换成原数据类型。
   • 当算子的计算过程比较繁琐时，可以通过抽调内部函数的方法保持每个模块的简洁性和可读性。