算子代码实现（TBE DSL）

简介

通过调用TBE DSL接口，在算子工程下的“tbe/impl/add.py”文件中进行Add算子的实现，包括算子函数定义、算子入参校验、compute过程实现及调度与编译。

代码模板介绍

MindStudio在“tbe/impl/add.py”中生成代码模板。

# 导入依赖的Python模块
import tbe.dsl as tbe
from tbe import tvm
from tbe.common.register import register_op_compute
from topi import generic


# 算子计算函数
@register_op_compute(("Add")
def add_compute(x, y, z, kernel_name="add"):
    """
    To do: Implement the operator by referring to the
           TBE Operator Development Guide.
    """

    res = tbe.XXX(x, y)
    return res


# 算子定义函数
def add(x, y, z, kernel_name="add"):
    """
    To do: Implement the operator by referring to the
           TBE Operator Development Guide.
    """
    # 输入参数占位
    data_x = tvm.placeholder(x.get("shape"), dtype=x.get("dtype"), name="data_x")
    data_y = tvm.placeholder(y.get("shape"), dtype=y.get("dtype"), name="data_y")

    # 调用算子计算函数
    res = add_compute(data_x, data_y, z, kernel_name)

    # 自动调度
    with tvm.target.cce():
        schedule = tbe.auto_schedule(res)

    # 编译
    config = {"name": kernel_name,
              "tensor_list": [data_x, data_y, res]}
    tbe.build(schedule, config)

引入算子开发时依赖的Python模块。
- “tbe.dsl”：引入TBE支持的特定域语言接口，包括常见的运算vmuls、vadds、matmul等。
  具体的接口定义可查看Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/python/site-packages/tbe/dsl目录下的Python函数。
- “tbe.tvm”：引入TVM后端代码生成机制。
  具体的接口定义可查看Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/python/site-packages/tbe/tvm目录下的Python函数，使用方法请参见https://docs.tvm.ai/。
- “tbe.common.register.register_op_compute”：提供了实现算子的UB自动融合的接口。
  具体的接口定义可查看Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/python/site-packages/tbe/common/register/register_api.py文件中的register_op_compute函数的定义。
模板生成以算子名称_compute命名的计算函数声明。
- 若创建算子工程时选择的“Sample Template”或者“Tensorflow Template”，输入输出参数及属性会根据原型定义自动生成。
- 若创建算子工程时选择“Empty Template”，默认生成的参数为一个输入与一个输出，不带属性。
模板生成以算子名称命名的定义函数的声明与部分实现，模板中提供的实现函数中的示例代码包含如下功能：
- 获取输入tensor的shape与dtype。
- 对输入参数进行校验。
- 对输入tensor进行占位。
- 调用算子的compute函数进行计算、调度与编译。

算子定义函数实现

开发者需要根据MindStudio生成的模板代码进行算子计算函数的实现，同时需要在算子定义函数中增加算子输入/输出/属性的校验代码。由于Add算子允许两个输入数据的shape不同，但算子计算接口tbe.dsl.vadd( )要求两个输入shape相同，因此需要对算子两个输入的shape进行广播，并对其进行校验。有助于在算子编译阶段，提前发现问题。修改后代码如下所示：

若开发者通过Windows上远程打开的MindStudio，复制以下代码时可能会出现无法复制的情况，解决方法请参见Windows上远程打开MindStudio时，复制的内容无法粘贴到编辑器窗口中。

      
       
         
         
           from __future__ import absolute_import
import tbe.dsl as tbe
from functools import reduce
from tbe import tvm
from tbe.common.register import register_op_compute
from tbe.common.utils import para_check
from tbe.common.utils import shape_util


SHAPE_SIZE_LIMIT = 2147483648

# 实现Add算子的计算逻辑
@register_op_compute("Add", op_mode="dynamic", support_fusion=True)
def add_compute(input_x, input_y, output_z, kernel_name="add"):
    # 将shape转换为list
    shape_x = shape_util.shape_to_list(input_x.shape)    
    shape_y = shape_util.shape_to_list(input_y.shape) 

    # shape_max取shape_x与shape_y的每个维度的大值
    shape_x, shape_y, shape_max = shape_util.broadcast_shapes(shape_x, shape_y,
                                                              param_name_input1="input_x",
                                                              param_name_input2="input_y")
    shape_size = reduce(lambda x, y: x * y, shape_max[:])
    if shape_size > SHAPE_SIZE_LIMIT:
        raise RuntimeError("the shape is too large to calculate")

    # 将input_x的shape广播为shape_max
    input_x = tbe.broadcast(input_x, shape_max)
    input_y = tbe.broadcast(input_y, shape_max)

    # 执行input_x + input_y
    res = tbe.vadd(input_x, input_y)

    return res

# 算子定义函数
@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT,
                            para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
def add(input_x, input_y, output_z, kernel_name="add"):
    # 获取算子输入tensor的shape与dtype
    shape_x = input_x.get("shape")
    shape_y = input_y.get("shape")

    # 检验算子输入类型
    check_tuple = ("float16", "float32", "int32")
    input_data_type = input_x.get("dtype").lower()
    para_check.check_dtype(input_data_type, check_tuple, param_name="input_x")

    # shape_max取shape_x与shape_y的每个维度的最大值
    shape_x, shape_y, shape_max = shape_util.broadcast_shapes(shape_x, shape_y,
                                                              param_name_input1="input_x",
                                                              param_name_input2="input_y")

     # 如果shape的长度等于1，就直接赋值，如果shape的长度不等于1，做切片，将最后一个维度舍弃（按照内存排布，最后一个维度为1与没有最后一个维度的数据排布相同，例如2*3=2*3*1，将最后一个为1的维度舍弃，可提升后续的调度效率）
    if shape_x[-1] == 1 and shape_y[-1] == 1 and shape_max[-1] == 1:
        shape_x = shape_x if len(shape_x) == 1 else shape_x[:-1]
        shape_y = shape_y if len(shape_y) == 1 else shape_y[:-1]
        shape_max = shape_max if len(shape_max) == 1 else shape_max[:-1]

    # 使用TVM的placeholder接口对输入tensor进行占位，返回一个tensor对象
    data_x = tvm.placeholder(shape_x, name="data_1", dtype=input_data_type)
    data_y = tvm.placeholder(shape_y, name="data_2", dtype=input_data_type)

    # 调用compute实现函数
    res = add_compute(data_x, data_y, output_z, kernel_name)

    # 自动调度
    with tvm.target.cce():
        schedule = tbe.auto_schedule(res)
    # 编译配置
    config = {"name": kernel_name,
              "tensor_list": (data_x, data_y, res)}
    tbe.build(schedule, config)

          

        

      
     

算子定义函数声明包含算子输入信息、输出信息以及内核名称。
```
def add(input_x, input_y, output_z, kernel_name="add"):
```
- input_x, input_y：Add算子的两个输入tensor，每个tensor需要采用字典的形式进行定义，包含shape和dtype等信息。输入tensor的个数需要与算子信息库定义“tbe/op_info_cfg/ai_core/add.ini”中的定义保持一致。
- output_z：输出tensor，包含shape和dtype等信息，字典格式，此字段为预留位。
  输出tensor的个数需要与算子信息库定义“tbe/op_info_cfg/ai_core/add.ini”中的定义保持一致。
- kernel_name：算子在内核中的名称（即生成的二进制文件以及算子描述文件的名称），用户自定义，保持唯一，只能是大小写字母、数字、“_”的组合，且必须是字母或者“_”开头，长度小于或等于200个字符。
算子输入校验和输出形状推导。
Add算子需要对两个输入tensor的shape进行校验，且仅支持数据类型float16, float32, int32，此外Add算子允许两个输入shape不同，因此需要调用shape_util.broadcast_shapes()生成广播后的shape并对其进行校验。
对输入tensor进行占位。
调用TVM的placeholder接口分别对两个输入tensor进行占位，并分别返回一个tensor对象。

5中的tensor_list的输入tensor需要是tvm.placeholder接口返回的tensor对象，所以此对象在后续计算过程实现中不能被替换。
调用add_compute计算函数。
```
res = add_compute(data_x, data_y, output_z, kernel_name)  
```
data_x、data_y为3生成的占位tensor对象。

计算函数的实现请参见Compute函数实现。
算子调度与编译实现。
调度配置config中的tensor_list为：
```
"tensor_list": (data_x, data_y, res)
```
分别为两个输入tensor与一个输出tensor。

Compute函数实现

开发者需要根据算子计算逻辑自定义实现算子compute函数，Add算子的Compute函数实现如下：

      
           @register_op_compute("Add", op_mode="dynamic", support_fusion=True)
def add_compute(input_x, input_y, output_z, kernel_name="add"):
    # 将shape转换为list
    shape_x = shape_util.shape_to_list(input_x.shape)    
    shape_y = shape_util.shape_to_list(input_y.shape) 

    # shape_max取shape_x与shape_y的每个维度的大值
    shape_x, shape_y, shape_max = shape_util.broadcast_shapes(shape_x, shape_y,
                                                              param_name_input1="input_x",
                                                              param_name_input2="input_y")
    shape_size = reduce(lambda x, y: x * y, shape_max[:])
    if shape_size > SHAPE_SIZE_LIMIT:
        raise RuntimeError("the shape is too large to calculate")

    # 将input_x的shape广播为shape_max
    input_x = tbe.broadcast(input_x, shape_max)
    input_y = tbe.broadcast(input_y, shape_max)

    # 执行input_x + input_y
    res = tbe.vadd(input_x, input_y)

    return res

add_compute函数的声明如下所示。
```
@register_op_compute("Add", op_mode="dynamic", support_fusion=True)
def add_compute(input_x, input_y, output_z, kernel_name="add")
```
装饰器@register_op_compute("Add", op_mode="dynamic", support_fusion=True)是DSL算子开发方式中必需的，其作用是整网运行时算子支持做UB自动融合，使得当前自定义算子可以在UB中根据UB融合规则自动与其他算子的计算进行组装。

其中：
- input_x，input_y：compute函数的入参，为在3中声明的输入tensor对应的placeholder，包含shape和dtype等信息。
- output_z：为1中算子接口函数透传过来的dict类型。
- kernel_name：算子在内核中的名称。
进行Add算子的计算逻辑的实现。
Add算子要求相加的两个tensor的shape相同，所以首先通过调用tbe.broadcast接口将两个输入tensor广播成相同的shape，然后调用tbe.vadd接口实现输入tensor的相加，并返回计算结果tensor。

算子编译验证

在算子python文件最下方添加main函数调用该算子，通过MindStudio编译算子实现文件，用于单算子代码的简单语法校验，代码示例如下所示：

        
             # 算子调用
if __name__ == '__main__':
    input_output_dict = {"shape": (5, 6, 7),"format": "ND","ori_shape": (5, 6, 7),"ori_format": "ND", "dtype": "float16"}
    add(input_output_dict, input_output_dict, input_output_dict, kernel_name="add")

右键单击“tbe/impl/add.py”，选择“Run 'add'”，编译算子。
如果编译没有报错，且在当前目录“tbe/impl”下生成kernel_meta文件夹，包括以下文件，则表示算子代码能够编译运行。
- 算子二进制文件*.o。
- 算子描述文件*.json：用于定义算子属性及运行时所需要的资源。

父主题： 算子开发过程