陈天奇的tvm更新了：VTA，开源AI芯片栈

陈天奇 发表于 tvm社区 量子位 搬运 | 公众号 QbitAI

华盛顿大学陈天奇博士近日在tvm社区介绍了VTA(Versatile Tensor Accelerator)项目，量子位全文搬运如下。

问题：不只是硬件设计

2016年的秋天，我收到我们系硬件的教授Luis的邮件邀请畅谈合作计划，当时正在酝酿TVM计划的我这样认识了同系的法国同学Thierry，开始了我们将近两年的合作。AI芯片和硬件加速是深度学习发展写下一个大的方向，而如何设计专用加速芯片无疑是大家都感兴趣的话题，Thierry也不例外。加速芯片专用性强，反而可以使得逻辑设计更加趋于简单，在设计思路上面也可以更加开放，不用受到传统指令集和功能完整性因素的制约，正好适合在研究前沿的高校同学一展身手。

然而，当我们的合作真正开始的时候，才发现问题比原来想的要困难的多：深度学习芯片的主要问题远远超过了硬件本身。深度学习硬件真正的目标，是希望大家都可以把自己的深度学习的模型跑在AI芯片上。而加速器带来的特殊定制导致对于每个加速器，我们都需要重新设计一套驱动，以及上面的软件栈—-包含编译，优化和部署模块。 虽然以高校研究人员的能力可以设计出不亚于工业级别的芯片设计，但是他还需要重头打造所有的这些模块，才可以让大家都真正的把自己的模型运行在他的加速器之上。

软件栈设计本身的巨大压力成为了让我们的项目完成的一道大大的鸿沟。而反观现在的行业现状，不论是高校还是公司，大家都面临的同样的一个问题 — 深度学习芯片从来不只是硬件的问题，而是涉及到一个从硬件，到驱动，到编译优化部署，再到模型本身的全栈问题。

一直以来我们不断在思考这个问题：如何降低AI芯片设计门槛，让每一个人都可以玩深度学习硬件加速和系统开发。

今天，由Thierry领衔设计，我合作参与的项目VTA(Versatile Tensor Accelerator)作为TVM stack的一部分正式亮相，也给出了对于这个问题的一个答案。

VTA是什么

VTA的全名叫做Versatile Tensor Accelerator，直译过来为灵活的张量加速器。首先，VTA是一个完全开源的深度学习加速器。但是VTA不光包含了加速器设计本身，完整的驱动，tvm编译的整合和直接从tvm前端python编译部署深度学习模型的完整开源工具链。本次发布的VTA包含了模拟器和FPGA部署模块。当然因为设计本身完全开源，有兴趣的同学也可以直接基于目前的设计进行流片设计。

作为tvm社区的一部分，VTA对于不同的人有着不同的意义：

对于硬件设计者而言，VTA就像一个硬件和软件栈的蓝图。为不论是基于VTA扩展定制，或者是给自己的加速器增加TVM的支持提供详细的参照。硬件设计者也可以通过tvm快速的直接在自己的硬件上运行深度学习程序，减少开放开销。VTA本身也会不断演化发布各种可定制设计来支持各种应用的需求。对于硬件设计研究者而言，VTA提供了一套从软件到硬件完全开源的测试平台。

对于深度学习的从业人员来说，新的深度学习优化往往需要考虑到加速硬件的特性。如何让系统支持未来的各种加速器优化成为了系统设计人员头疼的问题。VTA的设计总结了现在的深度学习加速器的特性，并且会不断更新反应最新的研究成果。系统设计者可以通过VTA为目标来完成优化深度学习加速器的一个基石。这也是TVM社区加速器支持发展的路线之一。同时完全开放的设计带来了更多软硬件结合设计的机会。对于深度学习算法同学来说，有一个开放的硬件平台可以让大家更好地研究如量化，低精度，压缩和设计针对加速硬件优化的模型。

设计细节: 精简硬件和软硬件协同设计

VTA的设计在同时发布的论文中有详细解说，麻雀虽小，五脏俱全，基本反应了TPU数据流风格的AI加速器设计经验。我就讲一些个人的体会。VTA的设计总结了目前已经有的深度学习芯片的设计经验。最后设计的结果是我和Thierry在两年之内不断讨论凝练的结果。

一开始Thierry给出的设计，比较偏向于复杂指令集(CISC)和固定的流水线 ，但我们发现RISC的微指令变成加上CISC访存和单元张量指令混合的风格可以做出更好的设计。原因非常简单，如果把碰到的问题都通过硬件的手段来解决，势必会给硬件增加很多负担。以深度学习应用为例，一般TPU类的加速器会有一个矩阵乘法核心，而卷积等操作都会转化到这个矩阵乘法核心来做。CISC的设计方式会使得我们的指令集里面直接有一个conv2d的卷积操作，然后由硬件来把这个指令解析成微指令。但是这么做并不是最好的，因为光是卷积，就需要支持stride, dilation，如果对于医学图像应用，我们需要支持三维卷积。因为细微的应用差别就需要重新改变硬件显然并不是我们想要的。这一点其实也体现在了深度学习加速器的学术研究上 — 大家都会说我们的硬件设计了XX功能解决了XX问题。比如我之前在会议上有看到专门用于处理deconvolution的GAN加速器。

但是事实上，并非所有的问题都适合在硬件解决。针对刚才的conv2d的例子，虽然深度学习芯片的核心是矩阵乘法或者卷积单元，但是如何利用这些简单的基本单元组合成我们想要的操作的微指令本身，其实可以直接暴露给软件来可编程。这样的设计让硬件“越来越笨”，而把更多的任务交给软件也来接。这正是RISC的思想。当然天下没有免费的午餐，更加灵活的设计需要更多的软件的支持，这也是VTA在tvm stack中协同设计的原因。

当然，深度学习加速器肯定和传统的RISC不同 —传统的CPU的指令为单个寄存器指令，一定程度上VTA的loop源语类似于比较灵活的表达loop的CISC。但是把大的任务打碎成基本的张量计算单元的这样的思想是一致的。这样的设计涉及到更多的软硬件协同，但是一旦处理的当，也可以使得加速器变得更加灵活。比如VTA目前的设计可以直接支持是谁三维卷积，对于之前的GAN加速器的例子也都可以直接通过编程实现。

VTA里面最让我喜欢的设计是显式的控制流依赖操作。利用这些控制流依赖我们可以直接在软件层面实现动态的流水线，达到掩盖内存开销的目的。有趣的是，当我们回过头来看这个过程时候，在VTA设计的这个过程我们不断重新发现了计算机体系结构教科书里面要解决的问题。比如内存读写掩盖的这个技术一定程度上可以看成是SMT(simultaneous multithreading)，我们现在的系主任是主要的发明人之一。为了用好这个设计，我们必须在软件层面做出创新，但是反过来使得VTA可以拥有更加灵活的流水线依赖。

除了控制流之外，我们还直接为VTA设计了一个完整的JIT驱动，并且可以通过RPC直接在自己笔记本上远程运行FPGA加速器实现，真正做到用python去编写加速器程序。这些设计无一包含了软件和硬件协同设计的痕迹，这也需要软件和硬件专家的同时努力。我相信未来的深度加速器也会往“更加聪明的软件，更加笨的硬件”，和更多软硬件协同方向发展。

未来展望

作为TVM社区加速器支持的重要一环，我相信VTA可以使得大家都可以玩加速器的目标更近一步。当然VTA的设计并非完美，我们也会不断的完善和扩展硬件和软件层，同时结合TVM社区的其他模块一起来做联合的优化。

围绕这个问题，这个方向还留下包括硬件抽象，协同设计，模型表示等许多非常有趣的问题等着大家去探索,也欢迎更多的小伙伴加入TVM社区。

资源链接

VTA英文博客： https://tvm.ai/2018/07/12/vta-release-announcement.html

技术文档： https://docs.tvm.ai/vta/

代码仓库： https://github.com/dmlc/tvm/tree/master/vta

有兴趣的小伙伴可以直接下载代码玩，目前支持FPGA运行。即使没有FPGA，也可以直接通模拟器运行编译出来的加速器代码，具体可以参考官方文档教程。

量子位 QbitAI · 头条号签约作者

վ'ᴗ' ի 追踪AI技术和产品新动态