NNVM compiler可以将前端框架中的工作负载直接编译到硬件后端,能在高层图中间表示(IR)中表示和优化普通的深度学习工作负载,也能为不同的硬件后端转换计算图、最小化内存占用、优化数据分布、融合计算模式。
编译器的典型工作流如下图所示:
这个编译器基于此前发布的TVM堆栈中的两个组件:NNVM用于计算图,TVM用于张量运算。
其中,NNVM的目标是将不同框架的工作负载表示为标准化计算图,然后将这些高级图转换为执行图。
TVM提供了一种独立于硬件的特定域语言,以简化张量索引层次中的运算符实现。另外,TVM还支持多线程、平铺、缓存等。
对框架和硬件的支持
编译器中的NNVM模块,支持下图所示的深度学习框架:
具体来说,MXNet的计算图能直接转换成NNVM图,对Keras计算图的直接支持也正在开发中。
同时,NNVM compiler还支持其他模型格式,比如说微软和Facebook前不久推出的ONNX,以及苹果CoreML。
通过支持ONNX,NNVM compiler支持Caffe2、PyTorch和CNTK框架;通过支持CoreML,这个编译器支持Caffe和Keras。
而编译器中的TVM模块,目前附带多个编码生成器,支持多种后端硬件,其中包括为X86和ARM架构的CPU生成LLVM IR,为各种GPU输出CUDA、OpenCL和Metal kernel。