華盛頓大學計算機系博士生陳天奇、以及上海交通大學和復旦大學的研究團隊提出一個基于學習的框架，以優化用于深度學習工作負載的張量程序。該研究使用基于機器學習的方法來自動優化張量運算核心并編譯AI工作負載，從而可以將最優的性能部署到所有硬件。實驗結果表明，該框架能夠為低功耗CPU，移動GPU和服務器級GPU提供與最先進手工調優庫相媲美的性能。

深度學習在我們的日常生活中已經無處不在。深度學習模型現在可以識別圖像，理解自然語言，玩游戲，以及自動化系統決策（例如設備放置和索引）。張量算符（tensor operators），如矩陣乘法和高維卷積，是深度學習模型的基本組成部分。

可擴展的學習系統依賴于手動優化的高性能張量操作庫，如cuDNN。這些庫針對較窄范圍的硬件進行了優化。為了優化張量算符，程序員需要從邏輯上等價的許多實現中進行選擇，但由于線程，內存重用， pipelining和其他硬件因素的不同，性能上的差別很大。

支持多種硬件后端需要巨大的工程努力。即使在當前支持的硬件上，深度學習框架和模型的開發也從根本上受到庫中優化操作符設置的限制，阻止了諸如操作符熔合（operator fusion）之類的優化，從而產生不受支持的操作符。

針對這個問題，華盛頓大學計算機系博士生陳天奇、以及上海交通大學和復旦大學的研究團隊提出一個基于學習的框架，以優化用于深度學習工作負載的張量程序（ tensor programs）。

摘要

我們提出一個基于學習的框架，以優化用于深度學習工作負載的張量程序（ tensor programs）。矩陣乘法和高維卷積等張量算符（ tensor operators）的高效實現是有效的深度學習系統的關鍵。然而，現有的系統依賴于手工優化的庫，如cuDNN，這些庫只有很少的服務器級GPU能很好地支持。對硬件有要求的操作庫的依賴限制了高級圖形優化的適用性，并且在部署到新的硬件目標時會產生巨大的工程成本。我們利用學習來消除這種工程負擔。我們學習了領域特定的統計成本模型，以指導在數十億可能的程序變體上搜索張量算符的實現。我們通過跨工作負載的有效模型遷移來進一步加快搜索速度。

實驗結果表明，我們的框架能夠為低功耗CPU，移動GPU和服務器級GPU提供與最先進手工調優庫相媲美的性能。

學習優化張量程序問題的形式化方法

我們提出以下問題：我們是否可以通過學習來減輕這種工程負擔，并自動優化給定硬件平臺的張量算符程序？本論文為這個問題提供了肯定的答案。我們建立了統計成本模型來預測給定的低級程序的程序運行時間。這些成本模型指導了對可能程序空間的探索。我們的成本模型使用可遷移的表示形式，可以在不同的工作負載之間進行泛化，以加速搜索。這一工作的貢獻如下：

我們提供了學習優化張量程序問題的一種形式化方法，并總結了其關鍵特征。

我們提出了一個基于機器學習的框架來解決這個新問題。

我們使用遷移學習將優化速度進一步提高2倍至10倍。

我們在這個框架中提供了詳細的組件設計選擇和實證分析。

在實際的深度學習工作負載的實驗結果表明，我們的框架提供的端到端性能改進比現有框架好1.2倍至3.8倍。

圖1：該問題的一個例子。 對于給定的張量算符規范 ，有多種可能的低級別程序實現，每種實現都有不同的loop順序， tiling 大小以及其他選項。每個選項都創建一個具有不同性能的邏輯等效程序。我們的問題是探索程序空間并找到一個優化的程序。

圖2：學習優化張量程序框架的概覽

學習優化張量程序算法

圖3：編碼低級別循環AST的可能方法的示例

表1：單batch的ResNet-18推理中所有conv2d操作符的配置。H，W表示高度和寬度，IC表示輸入通道，OC表示輸出通道，K表示 kernel大小，以及S表示stride大小。

討論和結論

我們提出了一種基于機器學習的框架來自動優化深度學習系統中張量算符的實現。我們的統計成本模型允許在工作負載之間進行有效的模型共享，并通過模型遷移加速優化過程。這個新方法的優秀實驗結果顯示了對深度學習部署的好處。

在我們的解決方案框架之外，這個新問題的具體特征使它成為相關領域創新的一個理想測試平臺，如神經程序建模、貝葉斯優化、遷移學習和強化學習。

在系統方面，學習優化張量程序可以使更多的融合操作符、數據布局和數據類型跨不同的硬件后端。這些改進對于改進深度學習系統至關重要。我們將開放我們的實驗框架，以鼓勵在這些方向進行更多的研究。