Ⅰ 深度學習推理引擎的一些思考
深度學習推理引擎，特別是走編譯優化路徑的引擎，是解決現代深度學習應用的關鍵技術之一。作為曾經的傳統編譯器開發者，我見證了從傳統編譯器到深度學習領域的轉變。在阿里巴巴工作期間，我積極推動基於編譯優化的深度學習推理引擎的開發，這一領域在短短幾年內已經取得了顯著進展。
在18年初加入阿里巴巴時，基於編譯優化的深度學習推理引擎並不佔據主流地位，當時流行的解決方案是TFLite和NCNN等框架，它們依賴高效GEMM加速庫或針對特定CPU的手寫匯編代碼。然而，我堅信編譯器優化才是解決這個問題的最佳途徑和未來方向。我的堅定信念得到了老闆的支持，我們順利推動了這個項目，最終取得了很好的落地效果。現今，TVM、GLOW、MLIR等框架已成為業界認可的編譯優化深度學習推理引擎。
我選擇TVM作為基礎框架，原因在於它在編譯優化領域與我的想法最為契合。盡管TVM在18年初並不成熟，但相比其他框架，它更加接近我的技術理想。雖然沒有從零開始開發，但我認為基於已有的成熟框架進行擴展，可以更快地解決問題，同時也節省了大量時間和資源。
做深度學習推理引擎，並非易事。一方面，它涉及矩陣計算的優化，這需要對編譯器、體系結構、高性能計算等領域有深刻理解。另一方面，它並非想像中那麼簡單，深度學習推理引擎的真正實現需要下足功夫，並涵蓋量化壓縮、圖優化、子圖分離、異構執行等多個方面。例如，在支持Hexagon DSP時，深入了解其體系結構與編程模型是關鍵。
對於深度學習推理引擎，接收模型是第一步，包括TensorFlow、TFLite、MXNet等。目前，這一環節在不同框架中都存在不足，模型支持程度參差不齊。以TVM為例，雖然已實現了部分運算元支持，但仍有大量運算元未覆蓋。對於開發者而言，可選擇在框架基礎上進行擴展，共同推動技術進步。
在解析模型後的計算圖階段，各框架實施的優化策略不盡相同，但普遍會涉及運算元融合等通用優化技術。TVM引入了類似LLVM的Pass機制，支持異構執行，但在自動化的異構子圖分離方面仍有改進空間。業界對於此問題的研究，包括基於傳統機器學習的Cost Model與強化學習的方法，為解決異構執行問題提供了新的思路。
計算圖對接外部推理引擎是另一個值得關注的點。在實際應用中，業務部門通常追求模型快速且高效執行。以GPU為例，TensorRT提供高性能推理能力，但可能不支持某些運算元。結合TVM與TensorRT，通過子圖分離實現不同運算元在不同硬體上的高效執行，是解決這一問題的有效途徑。這一策略同樣適用於NPU的支持，通過異構執行達到性能優化。
在運算元層面，卷積操作是優化的重點。編譯器、體系結構和高性能計算領域的專家在這方面能發揮關鍵作用。TVM的Auto TVM和Halide思想的應用，使得優化技術的移植和性能提升變得更為高效。編譯生成部分，TVM採用LLVM為基礎，簡化了代碼指令生成的工作量，同時提供了Tensorize機制，允許開發者自定義優化。
量化壓縮是深度學習領域的重要技術之一，TVM支持不同的量化策略，包括TFLite轉換和TensorRT量化。量化策略的選擇和優化是關鍵，業界的一些框架提供快速重訓練等工具，以提高模型准確度和執行效率。這方面的學習和實踐對於深度學習編譯器的發展具有重要意義。
總之，基於編譯優化的深度學習推理引擎是未來趨勢，而TVM等框架在這一領域展現出巨大潛力。隨著技術的不斷發展，我們可以期待更多創新和突破，為深度學習應用提供更高效、更靈活的解決方案。