神经网络构造中,算子组成了不同应用功能的网络结构。而张量加速引擎(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提供的TBE标准算子库和离线模型执行器。TBE标准算子库为昇腾AI处理器提供神经网络需要使用到的算子,离线模型执行器专门用来进行离线模型的加载和执行。运行管理器下层是驱动,与昇腾AI处理器进行底层交互。
如上图所示,运行管理器对外提供各种调用接口,如存储接口、设备接口、执行流接口、事件接口以及执行控制接口,不同的接口由运行管理引擎控制完成不同的功能。
存储接口提供设备上HBM(High BandwidthMemory,高带宽存储器)或DDR(Double Data Rate,双倍速率内存)内存的申请、释放和复制等,包括设备到主机、主机到设备以及设备到设备之间的数据拷贝。这些内存拷贝分为同步和异步两种方式:同步拷贝指内存拷贝完成后才能执行下一步操作,而异步拷贝指在拷贝的同时也可以同时执行其他操作。
设备接口提供底层设备的数量和属性查询,以及选中、复位等操作。在离线模型调用了设备接口后选中某特性设备后,则模型中所有的任务都将会在这个被选中的设备上执行。若执行过程中需要向其他设备派发任务,则需要再调用一次设备接口进行设备选中。
执行流接口提供执行流的创建、释放、优先级定义、回调函数设置、对事件的依赖定义和同步等,这些功能关系到执行流内部的任务执行,同时单个执行流内部的任务必须按顺序执行。
如果多个执行流之间需要进行同步,则需要调用事件接口,进行同步事件的创建、释放、记录和依赖定义等,确保多个执行流得以同步执行完成并输出模型最终结果。事件接口除了用于分配任务或执行流之间的依赖关系,还可以用于程序运行中的时间标记,记录执行时序。
执行时,还会用到执行控制接口,运行管理引擎通过执行控制接口和Mailbox完成核函数的加载和存储异步拷贝等任务的派发。
更多信息请关注昇腾开发者社区
作者:Torrenza