代码结构介绍

TBE DSL方式实现的算子代码结构如下所示：

# 导入依赖的Python模块
import te.lang.cce
from te import tvm
from te.platform.fusion_manager import fusion_manager
from topi import generic
from topi.cce import util

# 算子计算函数
# 装饰器函数fusion_manager.register可选，若算子实现逻辑中涉及reshape操作，不可使用此装饰器函数
@fusion_manager.register("add")
def add_compute(input_x, input_y, output_z, kernel_name="add"):
    """
    算子计算逻辑实现
    """
# 算子定义函数

def add(input_x, input_y, output_z, kernel_name="add"):
    """
    算子校验（可选）
    为输入tensor占位
    """

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

    # 自动调度
    with tvm.target.cce(): 
        schedule = generic.auto_schedule(res)        
    # 算子编译
    config = {"print_ir": False,
              "name": kernel_name,
              "tensor_list": (data_x, data_y, res)}
    te.lang.cce.cce_build_code(schedule, config)

# 可选，若实现此函数，且算子信息库中的input的dtype与format的dynamicFormat.flag配置为true，则可在算子融合阶段调用此函数实现dtype与format的推导。
def op_select_format(input_x, input_y, output_y, kernel_name="add"):
    ...
    ...
# 可选，若实现此函数，且算子信息库中的needCheckSupport的flag参数配置为true，则可在算子融合阶段调用此函数实现算子的dtype与shape的校验。
def check_supported(input_x, input_y, output_y, kernel_name="add"):
    ...
    ...

算子实现代码总体结构包含依赖Python模块的导入，算子定义函数实现，算子计算函数实现。

其中：

• 算子定义函数包含算子的校验，计算函数的调用以及调度与编译。
• 算子计算函数是对算子计算逻辑的实现。

下面详细介绍每个代码块的实现。