多面體編譯技術在軟硬協同設計中的應用

摘要
多面體模型是一項成熟的編譯最佳化技術,廣泛地應用在傳統的編譯器及AI編譯器中。本文關注的是將多面體技術應用到軟硬協同設計領域的空間探索(Design Space Exploration),針對Systolic Array的硬體架構設計多面體的分析模型,藉助於現有的多面體編譯技術在編譯的過程中實現對硬體架構的探索,達到最佳效能;也可以在特定硬體約束下,改變演算法設計適配硬體。
多面體模型介紹
多面體模型(Polyhedral)主要關注的是程式設計中的迴圈最佳化,兩層迴圈的迴圈變數的取值範圍可以構成一個平面,三層迴圈的迴圈變數可以組成一個長方體,如圖1所示,因此得名多面體模型。

圖1 不同迴圈層次的多面體表示[1]
多面體編譯最佳化關注的是在確保程式執行正確的前提下重組迴圈的結構,實現效能最佳。比如圖2的迴圈中,左圖表示的原始的二維迭代空間,藍色箭頭表示資料(黑點)之間存在依賴關係,對角線的綠色表示資料沒有依賴關係,經過變基操作之後變為右圖的表示式及迭代空間,從形狀看像是把多面體進行了變形,形象地體現出多面體最佳化的過程。當然,變形的目的是為了實現平行計算,達到更好的效能,具體分析可以參考[1][2]。

圖2 多面體模型的原始迴圈空間和變基空間[1]
脈動陣列(Systolic Array)介紹
脈動陣列是基於資料流的計算架構,具有十分規則的計算正規化,對於兩個大小為4×4的矩陣乘投射到一個二維的4×4的脈動陣列上,其計算架構和資料流動可以表示為下圖。本文的分析也針對二維脈動陣列展開。

圖3 脈動陣列的資料流表示
圖中的脈動陣列的特點是資料沿著水平和垂直方向流動,在計算單元(PE:Process Element)中完成乘加操作。
脈動陣列的空時轉換模型
針對脈動陣列的空時變換(Space-Time Transformation)的分析模型在上世紀九十年代已經提出[3]。
定義一個二維脈動陣列的大小為

,假設兩個矩陣A和B,大小分別為

,對於矩陣乘的操作表達為

,該公式的對應的虛擬碼為

圖4 矩陣乘的偽程式碼表示
基於等價計算的方法,將上述程式碼的表達轉換為自由賦值的符號表達(Assignment-Free Notation)
Eq<1>表示的是輸入引數,Eq<2>表示的是計算過程,Eq<3>表示的是輸出結果。
針對符號表達,可以將其投射到多面體模型中,

表示迭代的變數(Iteration Variables),他們組成了迭代空間(Iteration Domain),

每個變數(或稱之為Node)對應原始程式碼表現形式的一個計算實體(Instance)

。在上述第二個公式Eq<2>中

組成了迭代矩陣(Iteration Vector)。為了上述三個公式中的變數統一,將迭代變數k引入到Eq<1>和Eq<3>中。

脈衝陣列的計算可以從空間域和時間域兩個維度進行描述,空間域表示每個計算落在脈動陣列的計算單元的位置,時間域表示計算結果輸出的時刻。透過觀察,在圖3的脈動陣列中,第一個計算結果

的輸出時刻為t=4,最後一個計算結果的輸出的時刻是t=N1+N2+N3。N3可以看成是資料傳輸延遲,N1和N2是脈動陣列固有的延遲。

將前文提到的迭代矩陣

表示為向量

,定義一個時間向量(Time Vector)

,那麼計算實體

在脈動陣列上的發生的時刻t表示為:

為了計算PE的空間座標

,定義一個投影矩陣(Project Matrix)

,如下

對於一個計算實體

將會對映到具體的PE的座標是

,計算過程如下

另外,投影方向可以使用任何垂直於投影矩陣的單位向量

表示,其中

注意,投影矩陣和投影方向向量不是唯一的,此處是選擇了

。投影方向向量指定了具體的對映方向,且沿著該向量方向的節點投影到相同的PE上。

根據以上關於時間t和空間z的公式表達,那麼就可以將脈動陣列的計算過程透過空時變換(Space-Time Transformation) T描述如下,T為投影矩陣和時間向量的組合。
針對圖3的脈動陣列架構,空時轉換矩陣為
從公式表達上,可以看出

是一個N-1維的矩陣,描述了在計算實體對映到的PE的座標,

為一維向量,描述了每個計算實體在所投射的PE上的計算時間(執行順序)。

多面體模型的構建
本節的分析根據論文PolySA [4]展開,論文中多面體模型定義為

為迭代空間;迭代矩陣為

;散射函式(Scattering Function)

對應為上述的空時變換矩陣

,散射函式就是空間對映(Space Mapping)和時間排程(Timing Scheduling)的表達;訪問函式(Access Function)

,表示為對記憶體的訪問。

邏輯標籤(Logical Stamp)S表示的是空間域資訊和時間域資訊。
針對脈動陣列的空間探索,散射函式的不同決定了產生不同的脈動陣列架構,而散射函式是由投影矩陣P和時間向量

決定的,P由投影向量

決定(見上節的分析),因此空間探索的問題變為遍歷出符合要求的

根據脈動陣列結構的自身特點,一是每個邊上的計算單元存在一個週期的延遲(脈動陣列本身固有特點),因此一個邊上的計算時間都需要大於0,

表示位於邊上的迭代向量。

二是對映到一個PE的節點無法在同一時刻計算,即前面提到的投影方向向量和時間向量的乘積要大於0。
散射函式存在傳遞性,可以將其不斷投射下去。將多面體模型

定義為N維陣列

,根據散射公式轉化為N-1的陣列

,並類推下去。

比如從

投射到

,投射函式定義為

多面體

的邏輯標籤S表示為

脈動陣列空間探索例項

圖5 矩陣乘(I=J=K=2)的空間探索分析
圖5中的圖b和圖c是針對2×2的脈動陣列,圖d是將矩陣操作投射到1×2的脈動陣列上。
圖5中a表示的是多面體的散射函式和訪存函式,對於矩陣A,它被訪問的迭代變數是i或者k,因此它的訪存函式為
圖b表示的是沿著k軸

(此處採用

表示投影方向,和前文的

一樣)進行投射,那麼在i和j軸相同的迭代向量

將投射到PE00

將投射到PE11。和時間向量軸垂直的平面的迭代向量在同一時刻輸出結果,比如計算實體

同時在t=1時刻得到輸出結果。

圖c表示的也是K軸的對映,空間對映的結構和圖b相同,但時間向量選擇是

,計算實體

 ,

結果在同一時刻輸出。

圖d表示將圖c的空時資訊進行再一次的投射,改變脈動陣列的結構。新的投射矩陣變為了向量

,時間向量為

,投射方向為

構建了多面體的中間表示(IR)之後,可以利用現有的多面體編譯技術搭建編譯框架,如圖6所示。Design Optimizer是針對FPGA上的資源的最佳化,比如DSP,RAM,LUT等。最終產生HLS程式碼的配置引數,透過模板的方式生成HLS IP。

圖6 論文[4]PolySA的編譯框架
軟硬協同設計的可能性
1.上述方法將矩陣乘法操作透過多面體的表達,將其對映到脈動陣列的硬體架構上,透過控制不同的投射方向

和時間向量

,得到不同的空間對映和時間表達,二者也形成了空間探索的變數。在特定脈動陣列架構的約束下,可以獲得符合約束的投影方向向量

和時間向量

,如果引入具體的記憶體訪問模型,比如對記憶體層級(Memory Hierarchy)的描述,就可以很好地評估出效能,比如功耗,根據效能選擇合適的

,從而完成特定硬體架構下最優效能的探索。

2.滿足特定的演算法效能要求下,儘可能地降低硬體設計的開銷,比如降低脈動陣列的維度,可以透過建立多級的散射關係進行空間探索,將一種演算法實現投射到最小的硬體結構上。
3.將編譯技術引入到不同硬體需求的空間探索中。針對特定架構的硬體,如果追求效能最優,在編譯的環節調出效能最優的對映方法;如果追求功耗最低,那麼可以調出符合功耗要求的最佳對映方法;如果底層是可配置的硬體,那麼從軟體和硬體兩個維度調優,結合效能和功耗的雙重考慮。
結論與思考
本文討論了針對脈動陣列架構構建多面體模型的分析方法,並將多面體編譯技術應用到硬體架構空間的探索中。該方法的啟示在於將編譯技術融合到軟硬協同設計中,在編譯的過程中針對特定的需求進行軟硬體兩個維度的空間探索,達到協同調優。但是目前的多面體分析模型侷限於脈動陣列這種具有規則性的硬體架構上,尚且不能對映到所有的架構,這是需要後期研究的工作。
參考文獻
[1] Uday Bondhugula, Polyhedral Compilation Opportunities in MLIR
[2] 要術甲傑, Polyhedral Model—AI晶片軟硬體最佳化利器,
[3] Space-time transformation and systolic arrays, https://gyires.inf.unideb.hu/KMITT/a52/ch04s02.html
[4] Jason Cong and Jie Wang. PolySA: polyhedral-based systolic array auto-compilation. In ICCAD 2018.
關於壁仞科技研究院
壁仞科技研究院作為壁仞科技的前沿研究部門,旨在研究新型智慧計算系統的關鍵技術,重點關注新型架構,先進編譯技術和設計方法學,並將逐漸拓展研究方向,探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則,將積極投入各類產學研合作並參與開源社群的建設,為相關領域的技術進步做出自己的貢獻。
掃碼關注我們

相關文章