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。