[llvm] 010a8f7 - [CodeGen] Fix restore blocks' BasicBlock information in branch relaxation

[ZHU Zijia via llvm-commits]

Author: ZHU Zijia
Date: 2022-12-02T02:42:22+08:00
New Revision: 010a8f7a90bf6e6d77f07f0a02dd8a63dfeb32da

URL: https://github.com/llvm/llvm-project/commit/010a8f7a90bf6e6d77f07f0a02dd8a63dfeb32da
DIFF: https://github.com/llvm/llvm-project/commit/010a8f7a90bf6e6d77f07f0a02dd8a63dfeb32da.diff

LOG: [CodeGen] Fix restore blocks' BasicBlock information in branch relaxation

In branch relaxation pass, restore blocks are created and placed before
the jump destination if indirect branches are required. For example:

        foo
        sd      s11, 0(sp)
        jump    .restore, s11
        bar
        bar
        bar
        j       .dest
.restore:
        ld      s11, 0(sp)
.dest:
        baz

The BasicBlock information of the restore MachineBasicBlock should be
identical to the dest MachineBasicBlock.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D131863

Added: 
    

Modified: 
    llvm/lib/CodeGen/BranchRelaxation.cpp
    llvm/test/CodeGen/RISCV/branch-relaxation.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/BranchRelaxation.cpp b/llvm/lib/CodeGen/BranchRelaxation.cpp
index 87d5d053318f..b2431886a011 100644
--- a/llvm/lib/CodeGen/BranchRelaxation.cpp
+++ b/llvm/lib/CodeGen/BranchRelaxation.cpp
@@ -88,7 +88,9 @@ class BranchRelaxation : public MachineFunctionPass {
   bool relaxBranchInstructions();
   void scanFunction();
 
-  MachineBasicBlock *createNewBlockAfter(MachineBasicBlock &BB);
+  MachineBasicBlock *createNewBlockAfter(MachineBasicBlock &OrigMBB);
+  MachineBasicBlock *createNewBlockAfter(MachineBasicBlock &OrigMBB,
+                                         const BasicBlock *BB);
 
   MachineBasicBlock *splitBlockBeforeInstr(MachineInstr &MI,
                                            MachineBasicBlock *DestBB);
@@ -202,12 +204,20 @@ void BranchRelaxation::adjustBlockOffsets(MachineBasicBlock &Start) {
   }
 }
 
-/// Insert a new empty basic block and insert it after \BB
-MachineBasicBlock *BranchRelaxation::createNewBlockAfter(MachineBasicBlock &BB) {
+/// Insert a new empty MachineBasicBlock and insert it after \p OrigMBB
+MachineBasicBlock *
+BranchRelaxation::createNewBlockAfter(MachineBasicBlock &OrigBB) {
+  return createNewBlockAfter(OrigBB, OrigBB.getBasicBlock());
+}
+
+/// Insert a new empty MachineBasicBlock with \p BB as its BasicBlock
+/// and insert it after \p OrigMBB
+MachineBasicBlock *
+BranchRelaxation::createNewBlockAfter(MachineBasicBlock &OrigMBB,
+                                      const BasicBlock *BB) {
   // Create a new MBB for the code after the OrigBB.
-  MachineBasicBlock *NewBB =
-      MF->CreateMachineBasicBlock(BB.getBasicBlock());
-  MF->insert(++BB.getIterator(), NewBB);
+  MachineBasicBlock *NewBB = MF->CreateMachineBasicBlock(BB);
+  MF->insert(++OrigMBB.getIterator(), NewBB);
 
   // Insert an entry into BlockInfo to align it properly with the block numbers.
   BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(), BasicBlockInfo());
@@ -482,7 +492,8 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
   // Create the optional restore block and, initially, place it at the end of
   // function. That block will be placed later if it's used; otherwise, it will
   // be erased.
-  MachineBasicBlock *RestoreBB = createNewBlockAfter(MF->back());
+  MachineBasicBlock *RestoreBB = createNewBlockAfter(MF->back(),
+                                                     DestBB->getBasicBlock());
 
   TII->insertIndirectBranch(*BranchBB, *DestBB, *RestoreBB, DL,
                             DestOffset - SrcOffset, RS.get());

diff  --git a/llvm/test/CodeGen/RISCV/branch-relaxation.ll b/llvm/test/CodeGen/RISCV/branch-relaxation.ll
index 8afaf5e99cea..ba9b2d39c506 100644
--- a/llvm/test/CodeGen/RISCV/branch-relaxation.ll
+++ b/llvm/test/CodeGen/RISCV/branch-relaxation.ll
@@ -2788,14 +2788,14 @@ define void @relax_jal_spill_32_restore_block_correspondence() {
 ; CHECK-RV32-NEXT:    .zero 1048576
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB6_4
-; CHECK-RV32-NEXT:  .LBB6_8: # %dest_3
+; CHECK-RV32-NEXT:  .LBB6_8: # %dest_1
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB6_4: # %dest_1
 ; CHECK-RV32-NEXT:    #APP
 ; CHECK-RV32-NEXT:    # dest 1
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB6_5
-; CHECK-RV32-NEXT:  .LBB6_10: # %dest_3
+; CHECK-RV32-NEXT:  .LBB6_10: # %dest_2
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB6_5: # %dest_2
 ; CHECK-RV32-NEXT:    #APP