算子实现流程

基于TBE DSL进行算子开发时,由于算子调度依赖Auto Schedule机制,算子开发层面无法直接控制,因此开发者基于TBE DSL算子开发时只需要重点关注算子的计算过程的实现,同时在计算过程中可以通过一些方式进行算子精度和性能的提升。

基于TBE DSL方式实现算子的流程如图1所示。

图1 TBE DSL自定义算子实现流程图
  1. 算子分析:在开发算子代码之前需要分析算子的数学表达式,输入、输出以及计算逻辑的实现,明确需要调用的TBE DSL接口。

    若是使用的第三方开源网络,网络中有昇腾AI处理器不支持的算子,可以根据第三方开源网站中的算子实现进行分析,推理出算子的数学表达式。

    若TBE DSL接口无法满足算子实现的要求,请考虑使用TIK开发方式。

    算子分析请参见算子分析

  2. 计算实现:进行算子计算逻辑的实现,包含依赖python模块的导入,接口的声明,参数的初始化及校验、计算过程的描述。

    其中算子实现文件中一般包含两个模块:算子接口函数与compute函数,其中算子接口函数负责输入tensor的初始化及参数校验,compute函数用来描述算子的计算过程,并在算子接口函数中被调用。

    计算实现是TBE DSL开发的最重要部分,如何实现算子的功能、如何进一步对算子进行精度、性能上的优化,都需要在此阶段充分考虑。

  3. 算子调度:调用tbe.dsl.auto_schedule接口自动生成相关调度。
  4. 算子编译:调用tbe.dsl.build接口进行算子编译,最终生成算子目标文件.o与算子描述文件.json。