FAIR 開源 Tensor Comprehensions，讓機器學習與數學運算高性能銜接

知識 02-17

AI 研習社消息，Facebook AI 研究院於近日開源了 C++ 庫及數學語言 Tensor Comprehensions，它能有效填補研究人員於數學運算領域的溝通鴻溝，並基於各種硬體後端上大規模運行工程模型。它的最大亮點在於，它採用了 Just-In-Time 的編譯自動生成機器學習社區所需的高性能代碼。

也就是說，通過 Tensor Comprehensions，研究人員能夠以數學符號的方式進行編寫，系統能夠根據需求進行編譯調整，並輸出專業的代碼。

數量級增長

傳統意義上，如果要從零創造一個具有高性能表現的機器學習層需要兩個大步驟，這可能需要耗費工程師數天甚至數周的努力。

1.在 NumPy 層級，研究人員需要寫一個全新層，並在以 PyTorch 為代表的深度學習庫中鏈接已有運算，然後進行小規模測試。如果要運行大規模實驗，需要通過數量級加速檢測其代碼實現。

2.工程師隨後採用這個層，並為 GPU 和 CPU 撰寫有效代碼（但這需要滿足多個條件）：

1）工程師需要是一位在高性能計算頗有了解的專家，但這一人才目前非常稀缺；

2）文本、策略、寫代碼，debug，工程師需要樣樣精通；

3）將代碼與實際任務相關的後端相連接，如冗長的參數檢查和添加樣板集成代碼

這也直接導致近年來深度學習社區一直依賴以 CuBLAS, MKL, 和 CuDNN 為代表的高性能庫而構建運行於 GPU 和 CPU 的代碼。對於研究員而言，如何尋找一條新的實現思路成為了巨大挑戰。

而在 Tensor Comprehensions 中，研究員們得以將這一流程從幾天甚至幾周減少到數分鐘。這一開源包含了：

用簡單語法表達一系列機器學習概念的數學符號

基於 Halide IR 數學符號的 C ++前端

基於整數集庫（ISL）的 Just-in-Time 編譯器，

一個基於進化搜索的多線程、多 GPU 自動調節器

早期工作

Halide 是一種最近在高性能圖像處理領域頗受歡迎的語言，它採用類似的高級函數語法來描述一個圖像處理的 pipeline，隨後在單獨代碼塊中調度到硬體上，並且詳細到如何平鋪、矢量化、並行化和融合。對於具有專業知識的人而言，這是一種非常高效的語言；但對於機器學習從業者來說，這一難度並不小。Halide 的自動調度在研究上非常活躍，但對於 GPU 上運行的機器學習代碼，目前還沒有很好的解決方案。

Tensor Comprehension 將 Halide 編譯器作為調用庫。FAIR 研究員構建了 Halide 的中間表徵（IR）和分析工具，並與多面編譯進行技術配對，因此，用戶可以在無需理解運行原理的條件下使用相似的高級語法編寫層。此外，FAIR 研究員也找到了簡化語言的方法，不需要為縮減運算制定循環邊界。

細節實現

使用 Halide 和多面編譯技術，Tensor Comprehensions 能通過委託內存管理和同步功能以自動合成 CUDA 內核。這一編譯能夠針對特定尺寸對一般操作符進行融合，對快速本地內存、快速縮減和 Just-in-Time 專業化都能優化。因不嘗試涉足內存管理，因此這一流程能夠輕鬆集成到機器學習的任意框架，以及任何允許調用 C++ 函數的語言中。

與經典編譯器技術和庫方法相反的是，多面編譯能讓張量理解為每個新網路按需調度單個張量元素的計算。

在 CUDA 層面，它結合了仿射循環轉換、融合/分裂和自動並行處理，同時確保數據在存儲器層次結構中正確移動。

圖中數字表示最初計算張量元素的順序，箭頭表示它們之間的依賴關係。在這個例子中，數字旋轉對應於循環交換，能夠實現深度操作器的融合。

研究員提供了一個集成的多線程、多 GPU 自動調節庫，以推進搜索過程，它使用 Evolutionary Search 來生成和評估數千種實現方案，並選擇性能最佳的方案。只需調用 Tensor Comprehension 上的 tune 功能，用戶便能目睹性能的改善。最好的策略是通過 protobuf 序列化，在當下或離線情況下可重用。

在性能方面，Tensor Comprehensions 可能仍然有很大提升空間，但在條件滿足的情況下，它已經可以匹配或打破當前機器學習框架與手動調整庫集成的性能，並通過將代碼生成策略適應特定問題大小的能力來實現。下圖展現了 Tensor Comprehensions 與 Caffe2 和 ATen 的對比。