Exo 2：高效能運算的新程式語言，大幅減少程式碼量

許多企業為了打造支撐現代人工智慧系統的高效能函式庫，不惜重金聘請頂尖人才。例如，NVIDIA 開發了一些最先進的高效能運算（HPC）函式庫，建立了難以被超越的競爭壁壘。然而，如果幾名學生能在短短幾個月內，僅用幾百行程式碼就能與這些頂尖的 HPC 函式庫一較高下，那會是什麼情況？

這正是麻省理工學院電腦科學與人工智慧實驗室（CSAIL）的研究人員所展示的成果，他們開發了一款名為 Exo 2 的新程式語言。Exo 2 屬於一種新型程式語言，MIT 教授 Jonathan Ragan-Kelley 稱之為「使用者可排程語言」（USLs）。與傳統語言依賴編譯器自動生成最佳化程式碼不同，USLs 讓程式設計師掌握主導權，允許他們撰寫「排程」來明確控制編譯器如何生成程式碼。這使得效能工程師能夠將簡單的程式轉化為功能相同但速度大幅提升的複雜程式。

現有的 USLs（如初代 Exo）有一個限制，即其排程操作相對固定，這使得在不同「核心」（高效能函式庫中的獨立元件）之間重複使用排程程式碼變得困難。相比之下，Exo 2 允許使用者在編譯器外部定義新的排程操作，從而促進可重複使用的排程函式庫的建立。

該研究的主要作者、MIT 電機工程與電腦科學博士生兼 CSAIL 成員 Yuka Ikarashi 表示，Exo 2 可以將總排程程式碼量減少 100 倍，並在多種平臺上提供與頂尖實作相當的效能，包括支援許多機器學習應用的基礎線性代數子程式（BLAS）。這使得 Exo 2 成為專注於跨不同操作、資料型別和目標架構最佳化核心的 HPC 工程師的理想選擇。

「這是一種自下而上的自動化方法，而不是在高效能程式碼上進行機器學習或人工智慧搜尋，」Ikarashi 解釋道。「這意味著效能工程師和硬體實作者可以撰寫自己的排程函式庫，這是一組針對其硬體應用的最佳化技術，以達到峰值效能。」

Exo 2 的一大優勢在於，它透過在不同應用和硬體目標之間重複使用排程程式碼，大幅減少了所需的編碼工作量。研究人員在 Exo 2 中實作了一個約 2,000 行程式碼的排程函式庫，封裝了針對線性代數和特定硬體（如 AVX512、AVX2、Neon 和 Gemmini 硬體加速器）的可重複使用最佳化技術。該函式庫整合了超過 80 個高效能核心的排程工作，每個核心僅需十幾行程式碼，其效能可與 MKL、OpenBLAS、BLIS 和 Halide 相媲美，甚至更優。

Exo 2 還引入了一種名為「Cursors」的新機制，提供了所謂的「穩定參考」，用於在整個排程過程中指向物件程式碼。Ikarashi 指出，穩定參考對於將排程封裝在函式庫函式中至關重要，因為它使排程程式碼獨立於物件程式碼的轉換。

「我們認為 USLs 應該設計成可擴充套件的，而不是擁有一組固定的操作，」Ikarashi 表示。「透過這種方式，語言可以透過實作支援多樣化最佳化需求和應用領域的函式庫來發展，以支援大型專案。」

Exo 2 的設計讓效能工程師能夠專注於高階最佳化策略，同時確保底層物件程式碼在使用安全原語的情況下保持功能等效。未來，團隊希望擴充套件 Exo 2 對不同型別硬體加速器（如 GPU）的支援。此外，多個正在進行的專案旨在改進編譯器分析本身，包括正確性、編譯時間和表達能力。

這項研究已發表在 arXiv 預印本伺服器上。