Facebook發布張量理解庫，自動編譯高性能機器學習核心

夏乙 編譯整理

量子位 出品 | 公眾號 QbitAI

Facebook AI Research今天發布了張量理解（Tensor Comprehension），這是一個C ++庫，也是一種數學語言，它能夠自動、按需地及時編譯出機器學習所需的高性能代碼。

簡單來說，就是將（部分）人類能讀懂的高級編程語言代碼，轉換成CUDA等專為機器優化的代碼。

Facebook在博客中說，Tensor Comprehension有助於彌合研究人員與工程師之間的溝通鴻溝。研究人員往往專註於數學運算，而工程師專註於在多種硬體後端上運行大規模模型的實際需求。

以前，構建高性能機器學習層通常需要幾天到幾周的工作，分為兩個環節：

一、研究人員在numpy級別的抽象中寫一個新層，與PyTorch等深度學習庫鏈接起來，進行小規模測試。要把這些用來驗證想法的代碼用到大規模實驗上，需要加快一個數量級。

二、工程師根據這個層為GPU和CPU寫高效代碼。

這個過程中需要的工程師，得是高性能計算專家，這類人才數量非常有限。他們需要了解情境、制定策略、編寫並調試代碼，同時，這個將代碼移到後端的工作還涉及一些雜活，比如冗長的參數檢查、添加樣板集成代碼等。

因此，過去幾年裡，深度學習界在生成GPU和CPU尚運行的高性能代碼時，越來越依賴CuBLAS、MLK、CuDNN等代碼庫。

當想要試驗的新想法偏離了這些代碼庫提供的基元時，會涉及到大量的工程工作，研究人員會被嚇得止步不前。

Facebook認為，簡化從想法到高性能代碼的過程，把它從幾天、幾周縮短到幾分鐘有著重大的實踐價值。他們希望研究人員從數學的角度描述自己的想法之後，通過Tensor Comprehension的自動編譯和調整，就能生成性能良好的專用代碼。

這一版本的Tensor Comprehension包含：

能用簡單語法表達一系列機器學習想法的數學符號；

基於Halide IR的，表達這些數學符號的C++前端；

一個基於整數集庫（ISL）的多面即時（JIT）編譯器；

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

Tensor Comprehension將高性能圖像處理領域的流行語言Halide的編譯器作為一個庫，建立在Halide的中間表示（intermediate representation，簡稱IR）基礎上，和多面編譯（polyhedral compilation）技術結合起來。讓用戶可以用類似的高級語法編寫層，但不用說明它如何運行。

通過Halide IR和polyhedral compilation的結合，Tensor Comprehension通過委託內存管理和同步功能自動合成CUDA核。這種翻譯針對特定的運算符融合、快速本地內存、快速縮減和特定尺寸的JIT專門化進行優化。

Facebook在博客中表示，因為這一工作流不嘗試擁有或優化內存管理，因此可以輕鬆高效地集成到任何ML框架和能調用C++函數的語言中。

不同於傳統的編譯器技術和庫方法，多面編譯讓Tensor Comprehension能為每個新網路按需調度單個張量元素的計算。

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

圖中的數字表示張量元素最初計算的順序，箭頭表示它們之間的依賴關係，數字旋轉對應著循環交換，讓深層運算符融合成為可能。

Facebook還為推動搜索過程提供了一個集成的多線程、多GPU自動調整庫，用進化搜索來生成和評估數千種實現方案，並選擇性能最佳的方案。調用Tensor Comprehension上的調節（tune）功能就能實時看到性能的提升。

在性能方面，Tensor Comprehension在部分案例中可以匹敵甚至超越集成了手動調節代碼庫的機器學習框架。這主要依靠讓代碼生成策略適應特定問題大小的能力。

下圖可以看出，Tensor Comprehensions自動生成核的性能，比Caffe2和ATen用供應商提供的CuDNN等庫生成的核有所提升。

對於英偉達、英特爾等硬體廠商所提供的快速庫來說，Tensor Comprehension是一種補充，也可以和CuDNN、MKL、NNPack等一起使用。

稍後，Facebook還將發布Tensor Comprehension的PyTorch集成。

Tensor Comprehensions是Facebook，法國國家信息與自動化研究所（Inria），蘇黎世聯邦理工學院和麻省理工學院的合作項目。

論文：

Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, Albert Cohen

https://arxiv.org/abs/1802.04730

代碼：

https://github.com/facebookresearch/TensorComprehensions

文檔：

https://facebookresearch.github.io/TensorComprehensions/

博客：

https://research.fb.com/announcing-tensor-comprehensions/

—完—

加入社群

量子位AI社群13群開始招募啦，歡迎對AI感興趣的同學，加小助手微信qbitbot5入群；

此外，量子位專業細分群(自動駕駛、CV、NLP、機器學習等)正在招募，面向正在從事相關領域的工程師及研究人員。

進群請加小助手微信號qbitbot5，並務必備註相應群的關鍵詞~通過審核後我們將邀請進群。（專業群審核較嚴，敬請諒解）

誠摯招聘

量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復「招聘」兩個字。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！