TBE算子开发方式

昇腾AI软件栈提供了TBE(Tensor Boost Engine:张量加速引擎)算子开发框架,开发者可以基于此框架使用Python语言开发自定义算子,通过TBE进行算子开发有以下几种方式:

整个计算过程可以表现为多个输入张量经过一个计算节点得到多个张量的过程。TIK方式、DSL的开发流程本质上是一样的,只不过开发的抽象层次不一样而已。

表1 算子开发方式对比

参数

DSL方式

TIK方式

语言

Python

Python

运用场景

常用于各种算术逻辑简单向量运算,或内置支持的矩阵运算及池化运算,例如eltwise类操作。

适用各类算子的开发,对于无法通过lambda表达描述的复杂计算场景也有很好的支持,例如排序类操作。

入门难度

较低

较高

适用人群

入门开发者,需要了解NN、TBE DSL相关知识。

高级开发者,需要了解NN,昇腾AI处理器架构、TIK指令、数据搬运等相关知识。

特点

TBE DSL接口已高度封装,开发者仅需要使用DSL接口完成计算过程的表达,后续的Schedule创建、优化及编译都可通过已有接口一键式完成。

入门难度高,程序员直接使用TIK提供的API完成计算过程及Schedule过程,需要手工控制数据搬运的参数和Schedule。

开发者无需关注Buffer地址的分配及数据同步处理,由TIK工具进行管理。

不足

某些场景下性能可能较低,复杂算子逻辑无法支持表达。

TIK对数据的操作更加灵活,但需要手工控制数据搬运的参数和Schedule过程。

代码编写接近底层硬件架构,过程优化等基于特定硬件特性。