[ARM] Search backwards for CMP when combining into CBZ

[dmgreen]

Summary

The constant island pass currently only looks at the instruction immediately before a branch for a CMP to fold into a CBZ/CBNZ. This extends it to search backwards for the instruction that defines CPSR. We need to ensure that the register is not overridden between the CMP and the branch.

The test cases are new, I have just shown the differences for clarity.

[dmgreen]

Now with some mir tests.

[samparker]

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

1931

Do we need to still check this? I'm assuming we can't have a tCMPi8 with a high reg?

[dmgreen]

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

1931

Yeah, good point. It's only the two register CMP that can take higher registers. I'll removed it here.

[dmgreen]

Removed isARMLowRegister check, reversed the condition of DestOffset - BrOffset, and added some Kill flag handling and tests.

[samparker]

LGTM!

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

if (isCPEntryInRange(U.MI, UserOffset, U.CPEMI, MaxOffs, false, true)) {
       ++NumT2CPShrunk;
       MadeChange = true;
     }
   }
 
   return MadeChange;
 }
 
+static bool registerDefinedBetween(unsigned Reg,
+                                   MachineBasicBlock::iterator From,
+                                   MachineBasicBlock::iterator To,
+                                   const TargetRegisterInfo *TRI) {
+  for (auto I = From; I != To; ++I)
+    if (I->modifiesRegister(Reg, TRI))
+      return true;
+  return false;
+}
+
 bool ARMConstantIslands::optimizeThumb2Branches() {
   bool MadeChange = false;
 
   // The order in which branches appear in ImmBranches is approximately their
   // order within the function body. By visiting later branches first, we reduce
   // the distance between earlier forward branches and their targets, making it
   // more likely that the cbn?z optimization, which can only apply to forward
   // branches, will succeed.
else if (Pred == ARMCC::NE)
       NewOpc = ARM::tCBNZ;
     if (!NewOpc)
       continue;
     MachineBasicBlock *DestBB = Br.MI->getOperand(0).getMBB();
     // Check if the distance is within 126. Subtract starting offset by 2
     // because the cmp will be eliminated.
     unsigned BrOffset = getOffsetOf(Br.MI) + 4 - 2;
     unsigned DestOffset = BBInfo[DestBB->getNumber()].Offset;
-    if (BrOffset < DestOffset && (DestOffset - BrOffset) <= 126) {
-      MachineBasicBlock::iterator CmpMI = Br.MI;
-      if (CmpMI != Br.MI->getParent()->begin()) {
-        --CmpMI;
-        if (CmpMI->getOpcode() == ARM::tCMPi8) {
-          unsigned Reg = CmpMI->getOperand(0).getReg();
-          Pred = getInstrPredicate(*CmpMI, PredReg);
-          if (Pred == ARMCC::AL &&
-              CmpMI->getOperand(1).getImm() == 0 &&
-              isARMLowRegister(Reg)) {
-            MachineBasicBlock *MBB = Br.MI->getParent();
-            LLVM_DEBUG(dbgs() << "Fold: " << *CmpMI << " and: " << *Br.MI);
-            MachineInstr *NewBR =
-              BuildMI(*MBB, CmpMI, Br.MI->getDebugLoc(), TII->get(NewOpc))
-              .addReg(Reg).addMBB(DestBB,Br.MI->getOperand(0).getTargetFlags());
-            CmpMI->eraseFromParent();
-            Br.MI->eraseFromParent();
-            Br.MI = NewBR;
-            BBInfo[MBB->getNumber()].Size -= 2;
-            adjustBBOffsetsAfter(MBB);
-            ++NumCBZ;
-            MadeChange = true;
-          }
-        }
-      }
-    }
-  }
+    if (BrOffset >= DestOffset || (DestOffset - BrOffset) > 126)
+      continue;
+
+    // Search backwards to the instruction that defines CSPR
+    auto *TRI = STI->getRegisterInfo();
+    while (CmpMI != Br.MI->getParent()->begin()) {
+      if (CmpMI->modifiesRegister(ARM::CPSR, TRI))
+        break;
+    }
+
+    // Check that this inst is a CMP r[0-7], #0 and that the register
+    // is not redefined between the cmp and the br.
+    if (CmpMI->getOpcode() != ARM::tCMPi8)
+      continue;
+    if (Pred != ARMCC::AL || CmpMI->getOperand(1).getImm() != 0)
+      continue;
+    if (registerDefinedBetween(Reg, CmpMI->getNextNode(), Br.MI, TRI))
+      continue;
+
+    // Check for Kill flags on Reg. If they are present remove them and set kill
+    // on the new CBZ.
+    MachineBasicBlock::iterator KillMI = Br.MI;
+    bool RegKilled = false;
+    do {
+      --KillMI;
+      if (KillMI->killsRegister(Reg, TRI)) {
+        KillMI->clearRegisterKills(Reg, TRI);
+        RegKilled = true;
+        break;
+      }
+    } while (KillMI != CmpMI);
+
+    // Create the new CBZ/CBNZ
+        BuildMI(*MBB, Br.MI, Br.MI->getDebugLoc(), TII->get(NewOpc))
+            .addReg(Reg, getKillRegState(RegKilled))
+            .addMBB(DestBB, Br.MI->getOperand(0).getTargetFlags());
 
   return MadeChange;
 }
 
 static bool isSimpleIndexCalc(MachineInstr &I, unsigned EntryReg,
                               unsigned BaseReg) {
   if (I.getOpcode() != ARM::t2ADDrs)
     return false;
for (++J; &*J != JumpMI; ++J) {
     }
   }
 
   LLVM_DEBUG(dbgs() << "Removing Dead Add: " << *RemovableAdd);
   RemovableAdd->eraseFromParent();
   DeadSize += 4;
 }
 
-static bool registerDefinedBetween(unsigned Reg,
-                                   MachineBasicBlock::iterator From,
-                                   MachineBasicBlock::iterator To,
-                                   const TargetRegisterInfo *TRI) {
-  for (auto I = From; I != To; ++I)
-    if (I->modifiesRegister(Reg, TRI))
-      return true;
-  return false;
-}
-
 /// optimizeThumb2JumpTables - Use tbb / tbh instructions to generate smaller
 /// jumptables when it's possible.
 bool ARMConstantIslands::optimizeThumb2JumpTables() {
   bool MadeChange = false;
 
   // FIXME: After the tables are shrunk, can we get rid some of the
   // constantpool tables?
   MachineJumpTableInfo *MJTI = MF->getJumpTableInfo();

llvm/test/CodeGen/ARM/constant-island-cbz.mir

 liveins:
   - { reg: '$r0', virtual-reg: '' }
   - { reg: '$r1', virtual-reg: '' }
 body:             |
   ; CHECK-LABEL: name: test_notfirst
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.2(0x30000000), %bb.1(0x50000000)
   ; CHECK:   renamable $r0, $cpsr = tADDrr killed renamable $r0, renamable $r1, 14, $noreg
-  ; CHECK:   tCMPi8 renamable $r0, 0, 14, $noreg, implicit-def $cpsr
   ; CHECK:   renamable $r1 = t2ADDrs renamable $r0, killed renamable $r1, 18, 14, $noreg, $noreg
-  ; CHECK:   tBcc %bb.2, 0, killed $cpsr
+  ; CHECK:   tCBZ $r0, %bb.2
   ; CHECK: bb.1:
   ; CHECK:   renamable $r0 = tLDRi killed renamable $r0, 0, 14, $noreg :: (load 4 from %ir.x)
   ; CHECK:   tTAILJMPdND @c, 14, $noreg, implicit $sp, implicit $sp, implicit killed $r0
   ; CHECK: bb.2:
   ; CHECK:   $r0, dead $cpsr = tMOVi8 0, 14, $noreg
   ; CHECK:   tBX_RET 14, $noreg, implicit killed $r0
   bb.0:
     successors: %bb.1(0x30000000), %bb.2(0x50000000)

 tracksRegLiveness: true
 liveins:
   - { reg: '$r0', virtual-reg: '' }
   - { reg: '$r1', virtual-reg: '' }
 body:             |
   ; CHECK-LABEL: name: test_notredefined
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.2(0x30000000), %bb.1(0x50000000)
-  ; CHECK:   tCMPi8 renamable $r0, 0, 14, $noreg, implicit-def $cpsr
   ; CHECK:   renamable $r1 = t2ADDrs renamable $r0, killed renamable $r1, 18, 14, $noreg, $noreg
-  ; CHECK:   tBcc %bb.2, 0, killed $cpsr
+  ; CHECK:   tCBZ $r0, %bb.2
   ; CHECK: bb.1:
   ; CHECK:   renamable $r0 = tLDRi killed renamable $r0, 0, 14, $noreg :: (load 4 from %ir.x)
   ; CHECK:   tTAILJMPdND @c, 14, $noreg, implicit $sp, implicit $sp, implicit killed $r0
   ; CHECK: bb.2:
   ; CHECK:   $r0, dead $cpsr = tMOVi8 0, 14, $noreg
   ; CHECK:   tBX_RET 14, $noreg, implicit killed $r0
   bb.0:
     successors: %bb.1(0x30000000), %bb.2(0x50000000)

 tracksRegLiveness: true
 liveins:
   - { reg: '$r0', virtual-reg: '' }
   - { reg: '$r1', virtual-reg: '' }
 body:             |
   ; CHECK-LABEL: name: test_killflag_1
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.2(0x30000000), %bb.1(0x50000000)
-  ; CHECK:   tCMPi8 killed renamable $r1, 0, 14, $noreg, implicit-def $cpsr
   ; CHECK:   renamable $r0 = t2ADDrs killed renamable $r0, killed renamable $r0, 18, 14, $noreg, $noreg
-  ; CHECK:   tBcc %bb.2, 0, killed $cpsr
+  ; CHECK:   tCBZ killed $r1, %bb.2
   ; CHECK: bb.1:
   ; CHECK:   renamable $r0 = tLDRi killed renamable $r0, 0, 14, $noreg :: (load 4 from %ir.x)
   ; CHECK:   tTAILJMPdND @c, 14, $noreg, implicit $sp, implicit $sp, implicit killed $r0
   ; CHECK: bb.2:
   ; CHECK:   $r0, dead $cpsr = tMOVi8 0, 14, $noreg
   ; CHECK:   tBX_RET 14, $noreg, implicit killed $r0
   bb.0:
     successors: %bb.1(0x30000000), %bb.2(0x50000000)

 tracksRegLiveness: true
 liveins:
   - { reg: '$r0', virtual-reg: '' }
   - { reg: '$r1', virtual-reg: '' }
 body:             |
   ; CHECK-LABEL: name: test_killflag_2
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.2(0x30000000), %bb.1(0x50000000)
-  ; CHECK:   tCMPi8 renamable $r1, 0, 14, $noreg, implicit-def $cpsr
-  ; CHECK:   renamable $r0 = t2ADDrs killed renamable $r1, killed renamable $r0, 18, 14, $noreg, $noreg
-  ; CHECK:   tBcc %bb.2, 0, killed $cpsr
+  ; CHECK:   renamable $r0 = t2ADDrs renamable $r1, killed renamable $r0, 18, 14, $noreg, $noreg
+  ; CHECK:   tCBZ killed $r1, %bb.2
   ; CHECK: bb.1:
   ; CHECK:   renamable $r0 = tLDRi killed renamable $r0, 0, 14, $noreg :: (load 4 from %ir.x)
   ; CHECK:   tTAILJMPdND @c, 14, $noreg, implicit $sp, implicit $sp, implicit killed $r0
   ; CHECK: bb.2:
   ; CHECK:   $r0, dead $cpsr = tMOVi8 0, 14, $noreg
   ; CHECK:   tBX_RET 14, $noreg, implicit killed $r0
   bb.0:
     successors: %bb.1(0x30000000), %bb.2(0x50000000)

llvm/test/CodeGen/Thumb2/constant-islands-cbz.ll

 ;
 ; CHECK-T2-LABEL: test:
 ; CHECK-T2:       @ %bb.0: @ %entry
 ; CHECK-T2-NEXT:    .save {r4, lr}
 ; CHECK-T2-NEXT:    push {r4, lr}
 ; CHECK-T2-NEXT:    mov r4, r0
 ; CHECK-T2-NEXT:    movs r0, #0
 ; CHECK-T2-NEXT:    strd r0, r0, [r4, #4]
-; CHECK-T2-NEXT:    cmp r2, #0
 ; CHECK-T2-NEXT:    strd r0, r0, [r4, #12]
 ; CHECK-T2-NEXT:    mov r0, r4
-; CHECK-T2-NEXT:    beq .LBB0_2
+; CHECK-T2-NEXT:    cbz r2, .LBB0_2
 ; CHECK-T2-NEXT:  @ %bb.1: @ %if.else
 ; CHECK-T2-NEXT:    bl equeue_create_inplace
 ; CHECK-T2-NEXT:    mov r0, r4
 ; CHECK-T2-NEXT:    pop {r4, pc}
 ; CHECK-T2-NEXT:  .LBB0_2: @ %if.then
 ; CHECK-T2-NEXT:    bl equeue_create
 ; CHECK-T2-NEXT:    mov r0, r4
 ; CHECK-T2-NEXT:    pop {r4, pc}