整体架构
从开发者在不同深度学习开放平台上所撰写的代码,到最终在不同硬件平台上执行,这其中存在着不小的差距。深度学习框架更多关注开发者的便捷性,而硬件设备更注重能耗和效率。TVM/NNVM,定位于开源的深度学习编译器技术栈(open deep learning compiler stack),整体目的就是缩小这两者之间的差距。
需要注意的是,无论是深度学习框架还是硬件设备,两者都具有高度的多样性,倘若一对一的进行优化和编译,工作量会由于组合爆炸而激增。为此,TVM/NNVM借用了在CPU架构中常用的编程语言的思想,即用中间语言(IR)来实现不同语言之间的翻译。首先不同深度学习框架中的网络模型被NNVM翻译成为抽象的、硬件无关的计算图模型;之后,再通过TVM将上述计算图模型细化为与硬件关联的程序。
图是对TVM/NNVM整体架构的示意图。TVM技术堆栈由NNVM和TVM两部分组成。NNVM负责图层面的优化,TVM负责算子层面的优化。不同的深度学习平台,如MXNet,可以直接通过NNVM提供的接口转化为计算图;或可以通过CoreML、ONNX等深度学习标准格式进行转换。在这一过程中,NNVM同时为计算图添加层融合等图层面的优化技术。最终,NNVM生成与硬件无关的、优化的计算图,向下传递给TVM。TVM则负责为计算图中的每个操作如何在不同的平台上的执行添加高效地实现,并最终生成可以在不同硬件平台上可执行的代码。