TBE简介
TBE(Tensor Boost Engine)负责执行昇腾AI处理器中运行在AI Core上的算子,TBE提供了基于TVM(Tensor Virtual Machine)框架的自定义算子开发能力,通过TBE提供的API可以完成相应神经网络算子的开发。
首先看下什么是TVM,随着深度学习的广泛应用,大量的深度学习框架及深度学习硬件平台应运而生,但不同平台的神经网络模型难以在其他硬件平台便捷的运行,无法充分利用新平台的运算性能。TVM(Tensor Virtual Machine)的诞生解决了以上问题,它是一个开源深度学习编译栈,它通过统一的中间表达(Intermediate Representation)堆栈连接深度学习模型和后端硬件平台,通过统一的结构优化Schedule,可以支持CPU、GPU和特定的加速器平台和语言。TVM的架构详细介绍请参考https://tvm.apache.org/。
TBE的逻辑架构如图1所示。
TBE工具给用户提供了多层灵活的算子开发方式,用户可以根据对硬件的理解程度自由选择,利用工具的优化和代码生成能力,生成昇腾AI处理器的高性能可执行算子。
图1 TBE在软件栈中的逻辑架构图
- 前端框架:包含MindSpore以及第三方开源框架TensorFlow(TF:Google机器学习开源框架)、Caffe(快速特征嵌入的卷积结构)等。
- 图编译器(Graph Compiler):Graph Compiler是基于昇腾AI软件栈对不同的机器学习框架提供统一的IR(Intermediate Representation)接口,对接上层网络模型框架,例如TensorFlow、Caffe、PyTorch,Graph Compiler主要功能包括图准备、图拆分、图优化、图编译、图加载、图执行和图管理等(此处图指网络模型拓扑图),最终将用户输入的中间表达的计算图编译成可以在昇腾AI处理器运行的模型。
- 张量加速引擎(TBE):提供用户开发自定义算子所需工具。TBE通过IR定义为Graph Compiler的图推导提供必要的算子信息,通过算子信息库提供算子调用等信息,TBE生成的二进制程序可在昇腾AI处理器上执行。
TBE功能框架
TBE内部包含了算子逻辑描述模块、调度(Schedule)模块、中间表示(Intermediate Representation,IR)模块、编译优化(Pass)模块以及代码生成(CodeGen)模块如图2所示。
图2 TBE功能框架
- 算子逻辑描述:面向开发者,提供算子逻辑的编写的接口(Compute接口),使用接口来编写算子。
- 调度模块:用于描述指定shape下算子如何在昇腾AI处理器上进行切分,仍然使用的是TVM提供的调度原语进行描述。
- 中间表示模块:使用TVM的IR(Intermediate Representation)中间表示,包括IR变形、AST(Abstract Syntax Tree)的维护等功能。
- 编译优化(Pass):对生成的IR进行编译优化,优化的方式有双缓冲(Double Buffer)、流水线(Pipeline)同步、内存分配管理、指令映射、分块适配矩阵计算单元等。
- 代码生成模块(CodeGen):CodeGen生成类C代码的临时文件,这个临时代码文件可以通过编译器生成算子的实现文件,可被网络模型直接加载调用。