編譯器最佳化記錄(Mem2Reg+SSA Destruction)

Radioheading發表於2023-09-22

編譯器最佳化記錄(2) Mem2Reg+SSA Destruction

寫的時候忽然想起來,這部分的內容恰好是在我十八歲生日的前一天完成的。算是自己給自己的一份成長的紀念吧。

0. 哪些東西可以Mem2Reg

顧名思義,Mem2Reg的意思是我們可以透過維護每個函式中區域性變數被賦值之後產生的副本來消除對其alloca,而後進行一系列load/store的過程(眾所周知這一類操作是需要更多時間的)。

一般來說,所有的alloca都可以被消除,但是對於某個函式超過\(8\)個引數而言,情況會不太一樣。具體而言,它們是從棧上傳過來的,因而我們需要將其load到一個新的暫存器上。

另外,對於函式的\(0\to7\)號引數,在進行完Mem2Reg之後也有後續操作。我們需要在函式剛開始的時候把這幾個引數從物理暫存器\(a_0,...,a_7\)轉移到到我們給他們分配的虛擬暫存器上。

1. 確定插入phi的位置

在[上一份部落格](編譯器最佳化記錄(控制流圖,支配樹) - Radioheading - 部落格園 (cnblogs.com))中,我們已經獲得了插入phi的前置內容(即支配樹、支配邊界),那麼接下來我們就可以插入phi指令。

注意,下面的內容都是以函式為單位進行的,Mem2Reg的本質應該是一個Function Pass

根據上文所述的規則,我們可以找到每個可以被消除的區域性變數,只需要遍歷enterBlock中的各個alloca。接下來,我們對於每個可以消除的區域性變數\(alloca\),記錄它在每個塊中的最後一次\(def\),這樣的目的很明顯,是為了記錄在支配邊界,該變數應當被賦值為什麼。

我們採用HashMap<IRRegister, HashMap<BasicBlock, entity>> all來記錄對於每個區域性變數,它在每個塊中的最後一次定值。它肯定會發生改變,因為我們在某個塊中插入phi指令的時候,就又創造了一次定值。

然後,我們採用工作表演算法來解決。

工作表演算法(WorkList Algorithm)就是說,我們用一個集合\(W\)來作為工作表,每次選取其中的一個節點,進行一個操作,然後加入這個操作會影響的其他節點(假使它們不在其中),直到工作表為空。我們大部分的最佳化都會用到它。

我們使用工作表\(W\),它的初始值為所有對\(alloca\)進行賦值(即store)的集合。我們同時用HashSet<BasicBlock> F來記錄所有已經出現過關於alloca的接下來,我們選取其中的一個塊bb1,如果!all.get(alloca).containsKey(bb1),那就說明bb1中對這個變數的最後一個定值一定是phi指令,我們要對all進行更新。接下來,我們就需要列舉bb1的支配邊界來插入phi指令了。對於它的支配邊界中的一個塊Y,如果!F.contains(Y),那麼我們就需要插入phi指令啦。同時我們也要把Y分別加入FW中。

我對於phi指令的設計是這樣的:

public class IRPhi extends IRBaseInst {
    public HashMap<BasicBlock, entity> block_value = new HashMap<>();
    public IRRegister dest;
    public IRRegister origin;
}

同時,我在BasicBlock中加入了HashMap<IRRegister, IRPhi> phiMap來記錄一個塊中的所有phi,在執行toString()的時候優先輸出。

經過這個過程,我們可以獲得所有需要插入phi的地方。接下來,就是考慮從每條路徑來的時候,這個變數應該被賦值為什麼了。

2. 變數重新命名

一個直觀的想法是這樣的。對於每個區域性變數\(alloca\)和它出現的某個塊,如果這個塊中,在這條指令上面有對它的定值,那麼就直接使用這個定值定出來的東西。如果啥定值都沒有,那麼這肯定說明連phi都沒有,說明從控制流的角度來說,上一次對它定值一定是在它的直接支配節點或更上面。

那麼,我們為了尋找這個“最後一次定值”,就需要在支配樹上進行DFS。

接下來,我們考慮對於每個個塊的操作,也即visitBlock(BasicBlock block, Function func)

2.1 需要使用的資料結構

我們使用HashMap<IRRegister, entity> last_def來記錄當前每個變數最後\(def\)使用的值。例如%add = add i32 %0, 1; store i32 %add, ptr @s就可以被理解為,當前對變數s的最後一次\(def\)使用的是%add

我們使用HashMap<IRRegister, entity cur_name>來表示我們需要修改的虛擬暫存器。例如在上面兩條指令之後,%1 = load i32, ptr @s; %add1 = add i32 %1, 1中的%1就可以被替代為%add

需要注意的是,在進入到同一個塊的不同支配樹後繼時,last_def, cur_name都應該維持一致。這就需要我們每個塊內給它們開一個副本,在一個後繼訪問完後,把它們的副本重新賦回來,然後再訪問下一個後繼。

2.2 操作流程

首先自然是開副本,注意java的引用賦值特性

接下來,我們考察每個基本塊的所有指令(這裡指令包含所有的phi)。如果這條指令是一個store或者一個phi,那麼我們就修改last_def。如果這條指令是一個load,那麼我們就修改cur_name

然後我們調整後繼節點中phi的使用值。虛擬碼如下

 for (block的每個後繼succ) {
            for (succ 的每一個phi) {
                記element為這個phi原本的區域性變數
                if (last_def.get(element) != null) {
                    phi加入(block, last_def.get(element))的這個entry
                }
            }
        }

最後,我們訪問block的每一個支配後繼,並把副本賦回來,並刪掉和我們消除的這些區域性變數有關的load/store/alloca。

2.3 加入預設分支

考慮這樣一個情況,有基本塊\(BB_1, BB_2, BB_3\)滿足\(pred(BB_3) = \{BB_1, BB_2\}, pred(BB_1)=pred(BB_2)=enterBlock\)。如果我們一開始定義了一個區域性變數\(x\),並只在\(BB_1\)中對其定值,且在\(BB_3\)中使用之。那麼根據上面的操作,\(BB_3\)中關於\(x\)的phi指令只有一個源頭。然而,如果你嘗試用clang編譯它,會報一長串的錯誤。這是因為對於\(BB_3\)的前驅還包括\(BB_2\)。而規範應該是對於每一個前驅,都有一個賦值。於是我們需要對那些沒出現的前驅補充值,這裡姑且賦成初始值吧(i32, i8賦值為0ptr賦值為null)。

進一步思考,這其實算是對於原始碼的語義精化。換言之,在這裡我們可能改變了原始碼的意義(儘管它可能是不安全的)。

3. SSA Destruction

注意到,剛剛的插入phi的過程仍然保持了Single Static Assignment的性質。但是在之後指令選擇(指令選擇(instruction selection)是將中間語言轉換成彙編或機器程式碼的過程。在LLVM後端中具體表現為模式匹配)的階段,我們並沒有對phi指令的對應翻譯方法。那就需要我們在IR過渡到彙編的過程中把phi轉化成多次分別的賦值,這顯然會打破每個虛擬暫存器只能被定值一次的準則。

一個可以想見的轉化方法如下:

enter_main_0:
br label %for.cond_0

for.cond_0:
%i_phi_0 = phi i32 [ %inc_0, %for.inc_0 ], [ 0, %enter_main_0 ]
%x_phi_0 = phi i32 [ %add_0, %for.inc_0 ], [ 1, %enter_main_0 ]
%slt_0 = icmp slt i32 %i_phi_0, 10
br i1 %slt_0, label %for.body_0, label %for.end_0

for.inc_0:
%inc_0 = add i32 %i_phi_0, 1
br label %for.cond_0

for.body_0:
%add_0 = add i32 %x_phi_0, 1
br label %for.inc_0

for.end_0:
br label %exit_main_0

exit_main_0:
ret i32 %x_phi_0
}
enter_main_0:
%i_phi_tmp_0 = 0
%x_phi_tmp_0 = 1
br label %for.cond_0

for.cond_0:
%x_phi_0 = %x_phi_tmp_0
%i_phi_0 = %i_phi_tmp_0
%slt_0 = icmp slt i32 %i_phi_0, 10
br i1 %slt_0, label %for.body_0, label %for.end_0

for.inc_0:
%inc_0 = add i32 %i_phi_0, 1
%i_phi_tmp_0 = %inc_0
%x_phi_tmp_0 = %add_0
br label %for.cond_0

for.body_0:
%add_0 = add i32 %x_phi_0, 1
br label %for.inc_0

for.end_0:
br label %exit_main_0

exit_main_0:
ret i32 %x_phi_0
}

考慮%phi = phi i32 [0, bb1], [1, bb2]我們就在bb1/bb2的末端插入一個對%phi的賦值。這裡我用了一個並不存在的llvm-ir指令IRMove,並在指令選擇階段直接將其變成了Move。

當然,如果你仔細看了上面的這段程式碼,你會發現我是先把所有值賦給一個%tmp,再在出現phi的那個基本塊中將其賦值給%phi。這是為什麼呢?

我們可以回顧支配邊界的定義。可以想象,存在一種控制流圖使得某個節點的支配邊界有它自己。那這就會導致上面的方法對這兩種程式碼會執行不同的結果:

BB1:
%i_phi_0 = phi i32 [ %x_phi_0, %for.inc_0 ], [ 0, %enter_main_0 ]
%x_phi_0 = phi i32 [ %i_phi_0, %for.inc_0 ], [ 1, %enter_main_0 ]
...
BB1:
%x_phi_0 = phi i32 [ %i_phi_0, %for.inc_0 ], [ 1, %enter_main_0 ]
%i_phi_0 = phi i32 [ %x_phi_0, %for.inc_0 ], [ 0, %enter_main_0 ]
...
這兩個phi變數互相賦值,這樣它們的先後順序會影響它們在`%for.inc_0`這個塊中的賦值結果。為了解決之,我們採用了新增虛擬暫存器的策略。這某種程度上和時序邏輯有些相似。畢竟每個塊中的所有phi都應該是嚴格在同一時間完成的。

p.s.如果你讀過編譯器指導手冊的話,你會發現,我們的這個操作也省掉了增加空塊以避免資料競爭的操作。

修改完上述內容之後,你的Mem2Reg應該就能重新透過asm的所有測試點了。

4. 參考資料

[1] [SSA book](CnTransGroup/StaticSingleAssignmentBookChinese: 《Static Single Assignment Book》- 中文翻譯 (github.com))

[2] 編譯器指導手冊(預覽8.1)

相關文章