昇腾AI处理器软件栈--张量加速引擎（TBE）

神经网络构造中，算子组成了不同应用功能的网络结构。而张量加速引擎（Tensor Boost Engine）作为算子的兵工厂，为基于昇腾AI处理器运行的神经网络提供算子开发能力，用TBE语言编写的TBE算子来构建各种神经网络模型。同时，TBE对算子也提供了封装调用能力。在TBE中有一个优化过的神经网络TBE标准算子库，开发者可以直接利用标准算子库中的算子实现高性能的神经网络计算。除此之外，TBE也提供了TBE算子的融合能力，为神经网络的优化开辟一条独特的路径。

TBE提供了基于TVM开发自定义算子的能力，通过TBE语言和自定义算子编程开发界面可以完成相应神经网络算子的开发，TBE的结构如下图所示，包含了特性语言（Domain-Specific Language，DSL）模块、调度（Schedule）模块、中间表示（Intermediate Representation，IR）模块、编译器传递（Pass）模块以及代码生成（CodeGen）模块。

TBE算子开发分为计算逻辑编写和调度开发，其中特定域语言模块提供了算子计算逻辑的编写接口，直接基于特定域语言编写算子的计算过程和调度过程。算子计算过程描述指明算子的计算方法和步骤，而调度过程描述完成数据切块和数据流向的规划。算子每次计算都按照固定数据形状进行处理，这就需要提前针对在昇腾AI处理器中的不同计算单元上执行的算子进行数据形状切分，如矩阵计算单元、向量计算单元以及AI CPU上执行的算子对输入数据形状的需求各不相同。

在完成算子的基本实现过程定义后，需要启动调度模块中分块（Tiling）子模块，对算子中的数据按照调度描述进行切分，同时指定好数据的搬运流程，确保在硬件上的执行达到最优。除了数据形状切分之外，TBE的算子融合和优化能力也是由调度模块中的融合（Fusion）子模块提供的。

算子编写完成后，需要生成中间表示来进一步优化，而中间表示模块通过类似于TVM的IR格式来进行中间表示的生成。在中间表示生成后，需要将模块针对各种应用场景进行编译优化，优化的方式有双缓冲（Double Buffer）、流水线（Pipeline）同步、内存分配管理、指令映射、分块适配矩阵计算单元等。

在算子经过编译器传递模块处理后，由代码生成模块生成类C代码的临时文件，这个临时代码文件可以通过编译器生成算子的实现文件，并可以被离线模型执行器直接加载和执行。

综上所述，一个完整的自定义算子通过TBE中的子模块完成整个开发流程，从特定域语言模块提供算子计算逻辑和调度描述，构成算子原型后，由调度模块进行数据切分和算子融合，进入中间表示模块，生成算子的中间表示。编译器传递模块以中间表示进行内存分配等编译优化，最后由代码生成模块产生类C代码可供编译器直接编译。TBE在算子的定义过程不但完成了算子编写，而且还完成了相关的优化，提升了算子的执行性能。

TBE的三种应用场景如图所示。

• 一般情况下，通过深度学习框架中的标准算子实现的神经网络模型已经通过GPU或者其它类型神经网络芯片做过训练。如果将这个神经网络模型继续运行在昇腾AI处理器上时，希望尽量在不改变原始代码的前提下，在昇腾AI处理器上能发挥最大性能。因此TBE提供了一套完整的TBE算子加速库，库中的算子功能与神经网络中的常见标准算子保持了一一对应关系，并且由软件栈提供了编程接口供调用算子使用，为上层深度学习中各种框架或者应用提供了加速的同时尽量避免了开发昇腾AI处理器底层的适配代码。

• 如果在神经网络模型构造中，出现了新的算子，这时TBE中提供的标准算子库无法满足开发需求。此时需要通过TBE语言进行自定义算子开发，这种开发方式和GPU上利用CUDA C++的方式相似，可以实现更多功能的算子，灵活编写各种网络模型。编写完成的算子会交给编译器进行编译，最终执行在AI Core或AI CPU上发挥出芯片的加速能力。

• 在合适的场景下，TBE提供的算子融合能力会促进算子性能的提升，让神经网络算子可以基于不同层级的缓冲器进行多级别的缓存融合，使得昇腾AI处理器在执行融合后的算子时片上资源利用率获得显著提升。

综上，由于TBE在提供算子开发能力的同时也提供了标准算子调用以及算子融合优化的能力，使得昇腾AI处理器在实际的神经网络应用中，可以满足功能多样化的需求，构建网络的方法也会更加方便灵活，融合优化能力也会更好的提高运行性能。

更多信息请关注昇腾开发者社区

作者：Torrenza