推演一個Buffer分配的語法設計
問題:有一個向量處理器,內部有一片Buffer,大小為X,可以基於它進行平行計算,令操作符為opX,一個操作系列可以表述為opX(N)。
opX的運算元(無論輸入還是輸出),都在Buffer中,我們稱為一個Tensor(張量),Tensor表示Buffer中的首地址和長度。
opX之間沒有complete-start關係(簡稱NCS,None-Complete-Start),也就是說op1完成前,op2可以被投入執行,處理器有原語可以主動保證opX之間的依賴,這個問題不在本問題的考慮範圍內,我們只保證:“我們可以主動建立任何兩個opX之間的CS關係”。
現在的問題是:用什麼語法來表述Tensor的分配。
為了更容易理解問題,我們舉一個典型的場景,看看程式設計師要面對什麼問題大概是什麼樣的。比如程式要要完成一組計算,他的行為可能是這樣的:
LoadDataToTensor() Barrier() opX(N) Barrier() LoadDataToTensor() Barrier() opX(N)
但這樣的效率是很低的,因為在Load的時候opX的執行部件就會全部閒著。更好的辦法是把它們交叉起來:
loadTensor(Tensor1) loadTensor(Tensor2) Barrier(Tensor1) opX(N)(Tensor1) Barrier(Tensor1)#等opX(N)對Tensor1的處理完成 loadTensor(Tensor3) Barrier(Tensor2) opX(N)(Tensor2) Barrier(Tensor2) LoadTensor(Tensor4) ...
Tensor3、4明顯可以複用前面1、2的空間。
我們可以把這個行為歸結為這樣的語法:
for i in range(x) pipeline(depth=2): t = AllocTensor() loadTensor(t) opXWithBarrier(N)(t)
這樣我們可以構造一個深度為2的流水線,我們需要兩個t,發射完兩個序列後,我們等其中一個完成,然後我們讓t複用它的Buffer,然後才投入執行第三個迴圈。我們如果要支援這樣的語法,Tensor對Buffer的分配關係,應該提供什麼樣的語義才能提供最利於程式設計師選擇的程式設計介面?(當然,這只是其中一種典型場景)。
方案1:Alloc/Free
Tensor通過Alloc分配,Free釋放。分配後即可以被投入使用。這種方案明顯會導致Buffer利用不充分,Alloc, Free方式是用於“充裕記憶體”的場景的,很多時候容忍一定程度的浪費,作為處理器內部Buffer的分配演算法,顯然不適合,這個方案首先拋棄。
方案2:基於名稱空間的主動分配和釋放
我們把for裡面的opX(N)看做是一個名稱空間,Nested的for看做是一個子名稱空間,用Tensor(size)這個定義來實現一次Buffer分配,保證Tensor在離開它定義的名稱空間時自動被釋放。
這個方案本質上是方案1的收縮自由度後的變體,程式設計師不能對Buffer的利用進行控制,就會產生碎片,如前所述,對於處理器內部Buffer,碎片不可原諒。
這個方案其實還有另一個問題:基於Tensor來進行運算,依賴可以做在Tensor上,如果op1和op2使用相同的Tensor,我們就可以建立op1和op2之間的CS依賴。Buffer空間的分配和釋放本身並沒有Tensor依賴,那麼我們就不得不等待Nested的空間全部計算結束,否則就可以導致後續計算沒有空間繼續。這個讓程式設計師失去控制的機會,可能導致效能瓶頸。
方案3:半靜態Buffer規劃
上面那個流水線語義的核心是,要把Buffer根據流水線平分使用,為此,我們首先引入一個Tensor別名的設計,你可以把Tensor切出多塊出來單獨使用,比如你有一個Tensor:
Tensor A(shape=(2,3))
我們可以在這個Buffer的原位放另一個Tensor層疊上去:
Tensor A1(shape=3) alias A[0, :]
Tensor A2(shape=3) alias A[1, :]
A1和A2複用A的記憶體,但使用了另一個語義進行解釋。當我們進入一個流水線的時候,我們讓流水線來處理這個分配,這個語法就會變成這樣:
for i in [n, m] with t in Tensor1[0:1]: loadTensor(t, i) opXWithBarrier(N)(t)
這個語義表達的是:把Tensor1的第一維分成兩份,每份產生一個別名t,我們先完成兩次迴圈的發射,然後等第一個迴圈的結束,再投入下一次的發射。這樣,程式設計師始終知道整個Buffer的規劃,在迴圈內部他還可以把t再通過別人在分解過多個獨立的緩衝區,這樣,整個控制力就比較強了。
我暫時能想到的最優方案是這個方案3了,有人有更喜歡的應用模式嗎?