TBE简介

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在软件栈中的逻辑架构图

TBE功能框架

TBE内部包含了算子逻辑描述模块、调度(Schedule)模块、中间表示(Intermediate Representation,IR)模块、编译优化(Pass)模块以及代码生成(CodeGen)模块如图2所示。

图2 TBE功能框架