评分并提供意见反馈 :
华为采用机器翻译与人工审校相结合的方式将此文档翻译成不同语言,希望能帮助您更容易理解此文档的内容。 请注意:即使是最好的机器翻译,其准确度也不及专业翻译人员的水平。 华为对于翻译的准确性不承担任何责任,并建议您参考英文文档(已提供链接)。
代码结构介绍
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("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模块的导入,算子定义函数实现,算子计算函数实现。
其中:
- 算子定义函数包含算子的校验,计算函数的调用以及调度与编译。
- 算子计算函数是对算子计算逻辑的实现。
下面详细介绍每个代码块的实现。