TVMNN编译Compiler栈
内容纲要
前言
深度学习/神经网络应用日益广泛,多终端部署形成常态。从CPU、ARM、GPU到专用的神经网络加速器/深度学习处理器,不同的终端/不同的体系结构引起神经网络的碎片化;每一款设备特别是专用的加速芯片部署深度学习是一件费力不讨好的事情;同时近年来,虽然CNN加速器在学术上和产业上火热,但不同的研究人员/企业在CNN加速器上使用不同的指令集,不同的体系结构,没有统一的标准使开发人员在终端上部署。所以一款支持前端,各种后端,统一结构/指令集,扩展方便的NN编译栈在NN生态上的构建有着举足轻重的影响力,这就是TVM的魅力。
现在主流的深度学习训练框架是Caffe/PyTorch/TensorFlow/MxNet等,对CPU/CUDA支持得很好。如果想把训练好的神经网络部署到其它的终端设备,这就带了几个挑战:
主流框架不支持ARM/FPGA/ASIC
嵌入式终端不需要训练功能,对前向推理的速度有极大的要求
嵌入式终端性能/内存/存储有限,主流框架的臃肿不适合部署
终端指令集,架构没有统一标准,开发部署难度很大
这时需要一款软件编译栈,上接前端主流深度学习框架,后接各种终端;同时满足轻量级,高性能,高度扩展与灵活性,开发容易等要求。TVM 是深度学习系统的编译器堆栈。旨在缩小以算力为重点的深度学习框架与以性能和效率为重点的硬件后端之间的差距。TVM与深度学习框架协同工作,为不同的后端提供端到端编译。TVM支持主流的深度学习前端框架,包括TensorFlow, MXNet, PyTorch, Keras, CNTK;同时能够部署到宽泛的硬件后端,包括CPUs, server GPUs, mobile GPUs, and FPGA-based accelerators。
本文深度分析TVM的源代码(主要是FPGA-based accelerator/VTA方面),总结TVM的几个实现特点:
对接前端的神经网络模型配置文件,将前端网络转为自主设计的AST Graph/Relay IR;后端编译都基于这个Graph/Relay IR。这类似于实现一个通用的编译器,将C++/JAVA/Python等语言转换为自主的编译型语言IR;然后基于IR实现一个高性能的后端程序部署。
实现NN的Tensor Operator Libray:不同的后端采用不同的实现方式。
根据后端硬件将前端网络编译为Complied PackedFunc生成动态链接库,用于后端执行。
RunTime和Driver加载Complied PackedFunc调用硬件执行推理计算。
在HalideIR/LLVM框架/思想上实现编译系统。
前端的分离方式便于TVM兼容主流的深度学习框架,后端的编译运行分离使后端只需要部署一个轻量级的TVM(只包括RunTime和Driver)。这种思想能够兼顾通用性与性能。同时TVM在VTA(NN加速器)上的设计了一套通用的类RISC指令集,体系结构。TVM的整体思想与最终目标就是通用,形成一个深度学习全栈生态;但目前TVM实现的效果如何,让走近源代码分析一番。
调研目标
本文着重分析TVM的Relay IR 层,编译,VTA(NN加速器)源代码;ARM/CUDA部分目前不在源代码分析范围内。
TVM介绍
TVM是一个端到端的深度学习工具链,能够高效地把前端深度学习模型部署到CPUs、GPUs和专用的加速器上。TVM具有强大的兼容能力和扩展性,支持主流的深度学习前端框架,包括TensorFlow, MXNet, PyTorch, Keras, CNTK;同时能够部署到宽泛的硬件后端,包括CPUs, server GPUs, mobile GPUs, and FPGA-based accelerators。
TVM设计架构
TVM架构如下:
Tvm注意点:
TVM源码架构
TVM软件架构
TVM 源代码由三层构成,包括FrontEnd接口、Relay优化和BackEnd部署。目前着重研究BackEnd VTA层级和与硬件相关的Relay部分。设计上三层应该互不影响,各自独立;但源代码分析中,Relay层糅合了与硬件无关的图优化和与VTA相关的调度生成,没有分得很开。
FrontEnd
支持主流的深度学习前端框架,包括TensorFlow, MXNet, PyTorch, Keras, CNTK。目前TVM可以继承到PyTorch框架中优化、训练,而不是单纯地调用CNN模型接口。
Relay
根据具体硬件对原始计算图进行重构、张量优化、数据重排等图优化操作。源代码分析中,Relay层比较杂,干的事情比较多,既对接上层的图优化又对接硬件的调度器。
Relay及Tensorlization示例
BackEnd
后端支持ARM、CUDA/Metal/OpenCL及加速器VTA。
VTA实现原理及设计思想提炼
整体结构
VTA层次结构
VTA源代码架构有三个部分组成:
VTA Hardware
为了实现硬件的通用化计算,VTA硬件参考RISC指令集,按照Fetch—>Load—>Compute—>Store模式,将所有操作划分到这种粒度的计算,不设计专用复杂的计算模式。在牺牲一定性能的情况下,尽可能兼容更多的深度学习网络。
VTA硬件体系结构
指令集
VTA指令分为四大类:
VTA 指令结构
数据流
VTA relies on dependence FIFO queues between hardware modules to synchronize the execution of concurrent tasks. The figure below shows how a given hardware module can execute concurrently from its producer and consumer modules in a dataflow fashion through the use of dependence FIFO queues, and single-reader/single-writer SRAM buffers. Each module is connected to its consumer and producer via read-after-write (RAW) and write-after-read (WAR) dependence queues。
VTA依赖于硬件模块之间的依赖FIFO队列来同步并发任务的执行。下图显示了给定硬件模块如何通过使用依赖FIFO队列和单读写器SRAM缓冲区,以数据流方式从其生产者和消费者模块并发执行。每个模块通过读后写(RAW)和读后写(WAR)依赖队列连接到其使用者和生产者。
Dataflow and Dependency
The pseudo-code above describes how a module executes a given instruction predicated on dependences with other instructions. First, the dependence flags within each instruction are decoded in hardware. If the instruction has an incoming RAW dependences, execution is predicated upon receiving a RAW dependence token from the producer module. Similarly, if the task has an incoming WAR dependence, execution is predicated upon receiving a WAR dependence token from the consumer module. Finally when the task is done, we check for outgoing RAW and WAR dependences, and notify the consumer and producer modules respectively.
上面的伪代码描述了一个模块如何根据与其它指令的依赖关系,执行给定的指令。首先,在硬件中对每条指令中的依赖标志进行解码。如果指令具有传入的原始依赖项,则在从生产者模块接收到原始依赖项令牌时,将断言执行。类似地,如果任务具有传入的WAR依赖性,则在从使用者模块接收到WAR依赖性令牌时,就断言执行。最后,当任务完成时,检查输出的RAW和WAR依赖,并分别通知使用者和生产者模块。
控制流
按照Fetch—>Load—>Compute—>Store模式去计算:
Inp buffer和Out buffer之间没有数据交换,固定buffer in/out。控制逻辑会处理数据/指令依赖以及memory latency hiding。
VTA Config
VTA主要对硬件的数据位宽,SRAM 大小,计算阵列大小进行配置,而不能更改大的计算架构,数据流,控制流等;同时TVM根据已配置好的VTA执行编译工作,而不会在编译阶段在线生成硬件代码。
VTA配置
Pyng HLS
VTA采用高层次综合C++实现Xilinix Pynq FPGA的部署,包含配置文件,Vivado HLS 脚本,HLS C++硬件实现三部分组成。
Pynq FPGA部署
Hw_spec和VTA Implementation构成HLS 实现的主体:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// GEMM Layout
// _____________________________|_type______________|
// arg 0: opcode | opcode_T |
// arg 1: pop_prev_dependence | bool |
// arg 2: pop_next_dependence | bool |
// arg 3: push_prev_dependence | bool |
// arg 4: push_next_dependence | bool |
// arg 5: reset_reg | bool |
// arg 6: uop_bgn | uop_idx_T |
// arg 7: uop_end | uop_idx_T |
// arg 8: iteration count ax0 | loop_T |
// arg 9: iteration count ax1 | loop_T |
// arg a: accum idx factor ax0 | acc_idx_T |
// arg b: accum idx factor ax1 | acc_idx_T |
// arg c: input idx factor ax0 | inp_idx_T |
// arg d: input idx factor ax1 | inp_idx_T |
// arg e: weight idx factor ax0 | wgt_idx_T |
// arg f: weight idx factor ax1 | wgt_idx_T
VTA HLS C++ 实现
硬件设计思想提炼
VTA的核心思想,将计算划分到一个通用的细腻度的计算结构Operation Wrapper:
Load Inst:加载Operation对应指令
Decode:译码
AGU:计算Operation所需要的DRAM或SRAM地址
Operation:利用Operation对应的硬件计算资源执行操作
无论是Fetch、load、compute、store等较大粒度的操作都可以利用OperationWrapper分解到更细粒度的四个步骤,这和传统的设计方法不同。举例,传统设计是将芯片分为DMA—>Load—>Compute—>Store四个模式,每个模式单独去设计一套尽量通用的操作单元,去匹配软件算法;同时每个模式不会再分解到一个更通用的细腻度结构;OperationWrapper将Fetch—>Load—>Compute—>Store中每个模式都分解为Load—>Decode—>AGU—>Operation,即使是传统认为Load/Store已经是最小粒度,TVM依然再次分解到OperationWrapper。这有点像RISC中再RISC的思想。
VTA 硬件设计思想
另外,传统方法设计一个统一的AGU负责所有的地址计算;而TVM通过OperationWrapper将统一的AGU分解到每个模块,让每个模块负责一个更细粒度的AGU。这种方式能够加强AGU的通用性。
VTA这种OperationWrapper的思想,将硬件进一步分解到更加通用的细粒度模式,比起传统方法通用型更强。
VTA的第二个核心思想,利用Stream传输并通过同步控制解决指令与数据的依赖性(有点类似操作系统中的设计理念),让CPU与VTA的通讯、传输、控制变得更向操作系统级靠拢,可能也是为了TVM的通用性考虑。
参考文献
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
VTA: A Hardware-Software Blueprint for Flexible Deep Learning Specialization
手机扫一扫
移动阅读更方便
你可能感兴趣的文章