[AArch64] Optimize compare and branch sequence when the compare's constant operand is power of 2

[bmakam]

Summary

Peephole optimization that generates a single TBZ/TBNZ instruction
for test and branch sequences like in the example below. This handles
the cases that miss folding of AND into TBZ/TBNZ during ISelLowering of BR_CC

Examples:

and  w8, w8, #0x400
cbnz w8, L1

to

tbnz w8, #10, L1

[t.p.northover]

include/llvm/Target/TargetInstrInfo.h

1109 ↗

(On Diff #50001)

I think you're misinterpreting optimizeCondBranch it's for any conditional control flow at the end of a basic block, and in fact the AArch64 version already handles some TBZ cases.

So I think the logic of this function should be folded into the other (as neatly as possible).

[bmakam]

Address Tim's comment.

[bmakam]

Thanks for the review Tim,

I have folded the logic into optimizeCondBranch as you suggested. Does this look reasonable?

[mcrosier]

Please run clang-format on your patch, Balaram.

lib/Target/AArch64/AArch64InstrInfo.cpp

2956

s/test/compare/

I. e., Replace compare and branch sequence by..?

3015

Please add a period. It might also be helpful to comment your converting this sequence to tbz/tbnz when the AND is a constant power of 2.

3029

Might be easier to read as: Opc != ANDW && Opc != ANDX. This should also remove the unneeded parens.

3052

Should you be checking the Imm value here? Isn't it possible to have an Imm <= 31, but still be defining a 64bit register?

I'm wondering if the condition should be something more like (AndDefMI->getOpcode() == AArch64::ANDXri), as is done above?

3054

Single statement; No need for curly brackets.

You could also use a select statement.

3061

Same.

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.cpp

3016

This looks overly specialized to me: "look through precisely one COPY, but don't accept a direct AND".

Shouldn't we make the look-through-copy logic generic (so it applies to the CSINC/CBZ too), and iterative?

3073

I'd be inclined to make this whole thing a switch on DefMI->getOpcode() once it's been found. They're handled more like sibling cases anyway, and I could easily see us adding other patterns later.

[bmakam]

Thanks for the comments Chad and Tim,
I will post an updated patch including all your suggestions.

lib/Target/AArch64/AArch64InstrInfo.cpp

3016

Thanks Tim, this looks reasonable to me. I will post an updated patch including this suggestion.

3052

The immediate is the bit that is tested whether it is set or not, so if it is less than 32 and x[0-9]+, #c is same as w[0-9]+, #c. Furthermore the integrated linker treats and x[0-9]+, #c differently and transforms it into and x[0-9]+, 32+#c if c <32. I think this is a bug in the integrated linker, if we use no-integrated-linker and x[0-9]+, #c is transformed into and w[0-9]+, #c automatically. I therefore am checking the immeadiate value to avoid dependency on the underlying linker.

3073

Sounds reasonable to me again, will do in updated patch. Just out of curiosity, what other patterns do you think can be added?

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.cpp

3073

The one I thought of the other day was MOV/CBZ -> B. "ORR #imm" too. I've no idea if these occur enough to be worthwhile though, so no need to add them unless you actually want to.

[bmakam]

Address review comments.

[t.p.northover]

Thanks for updating the patch. This looks good to go now.

Tim.

[chapuni]

llvm/trunk/lib/Target/AArch64/AArch64InstrInfo.cpp

2980 ↗

(On Diff #50295)

2nd "\brief" should not be here. [-Wdocumentation]

[mcrosier]

llvm/trunk/lib/Target/AArch64/AArch64InstrInfo.cpp

2980 ↗

(On Diff #50295)

Committed r263942.

lib/Target/AArch64/AArch64InstrInfo.cpp

 ///   b.<inverted condition code>
 ///
 /// 2.
 ///   csinc w9, wzr, wzr, <condition code>
 ///   tbz   w9, #0, 0x44
 /// to
 ///   b.<condition code>
 ///
+/// \brief Replace test and branch sequence by TBZ/TBNZ instruction when
+/// the test's constant operand is power of 2.
+///
+/// Examples:
+///   and  w8, w8, #0x400
+///   cbnz w8, L1
+/// to
+///   tbnz w8, #10, L1
+///
 /// \param  MI Conditional Branch
 /// \return True when the simple conditional branch is generated
 ///
 bool AArch64InstrInfo::optimizeCondBranch(MachineInstr *MI) const {
   bool IsNegativeBranch = false;
   bool IsTestAndBranch = false;
   unsigned TargetBBInMI = 0;
   switch (MI->getOpcode()) {
bool AArch64InstrInfo::optimizeCondBranch(MachineInstr *MI) const {
   MachineFunction *MF = MBB->getParent();
   MachineRegisterInfo *MRI = &MF->getRegInfo();
   unsigned VReg = MI->getOperand(0).getReg();
   if (!TargetRegisterInfo::isVirtualRegister(VReg))
     return false;
 
   MachineInstr *DefMI = MRI->getVRegDef(VReg);
 
+  // Look for AND if MI is CBZ/CBNZ and DefMI is a COPY instruction
+  if (!IsTestAndBranch && DefMI->isCopy()) {
+    if (DefMI->getParent() != MBB) return false;
+
+    if (!MRI->hasOneNonDBGUse(VReg)) return false;
+
+    unsigned AndVReg = DefMI->getOperand(1).getReg();
+    if (!TargetRegisterInfo::isVirtualRegister(AndVReg))
+      return false;
+
+    if (!MRI->hasOneNonDBGUse(AndVReg)) return false;
+
+    MachineInstr *AndDefMI = MRI->getVRegDef(AndVReg);
+    if (!AndDefMI) return false;
+    if (!(AndDefMI->getOpcode() == AArch64::ANDWri) &&
+        !(AndDefMI->getOpcode() == AArch64::ANDXri))
+      return false;
+
+    uint64_t Mask = AArch64_AM::decodeLogicalImmediate(
+                    AndDefMI->getOperand(2).getImm(),
+                    (AndDefMI->getOpcode() == AArch64::ANDXri) ? 64 : 32);
+    if (!isPowerOf2_64(Mask))
+      return false;
+
+    MachineOperand &MO = AndDefMI->getOperand(1);
+    unsigned NewReg = MO.getReg();
+    if (!TargetRegisterInfo::isVirtualRegister(NewReg))
+      return false;
+
+    assert(!MRI->def_empty(NewReg) && "Register must be defined.");
+
+    MachineBasicBlock &RefToMBB = *MBB;
+    MachineBasicBlock *TBB = MI->getOperand(1).getMBB();
+    DebugLoc DL = MI->getDebugLoc();
+    unsigned Imm = Log2_64(Mask);
+    unsigned Opc;
+
+    if (Imm > 31) {
+      MRI->constrainRegClass(NewReg, &AArch64::GPR64RegClass);
+      if (IsNegativeBranch) {
+        Opc = AArch64::TBNZX;
+      } else {
+        Opc = AArch64::TBZX;
+      }
+    } else {
+      MRI->constrainRegClass(NewReg, &AArch64::GPR32RegClass);
+      if (IsNegativeBranch) {
+        Opc = AArch64::TBNZW;
+      } else {
+        Opc = AArch64::TBZW;
+      }
+    }
+    BuildMI(RefToMBB, MI, DL, get(Opc)).addReg(NewReg).addImm(Imm).addMBB(TBB);
+    MI->eraseFromParent();
+    return true;
+  }
+
   // Look for CSINC
   if (!(DefMI->getOpcode() == AArch64::CSINCWr &&
         DefMI->getOperand(1).getReg() == AArch64::WZR &&
         DefMI->getOperand(2).getReg() == AArch64::WZR) &&
       !(DefMI->getOpcode() == AArch64::CSINCXr &&
         DefMI->getOperand(1).getReg() == AArch64::XZR &&
         DefMI->getOperand(2).getReg() == AArch64::XZR))
     return false;
 
   if (DefMI->findRegisterDefOperandIdx(AArch64::NZCV, true) != -1)

test/CodeGen/AArch64/aarch64-tbz.ll

+; RUN: llc -mtriple=aarch64-linux-gnueabi < %s | FileCheck %s
+
+; CHECK-LABEL: test1
+; CHECK: tbz {{w[0-9]}}, #3, {{.LBB0_3}}
+; CHECK: tbz [[REG1:x[0-9]+]], #2, {{.LBB0_3}}
+; CHECK-NOT: and [[REG2:x[0-9]+]], [[REG1]], #0x4
+; CHECK-NEXT-NOT: cbz [[REG2]], {{.LBB0_3}}
+
+; CHECK: b
+define void @test1(i64 %A, i64 %B) {
+entry:
+  %and = and i64 %A, 4
+  %notlhs = icmp eq i64 %and, 0
+  %and.1 = and i64 %B, 8
+  %0 = icmp eq i64 %and.1, 0
+  %1 = or i1 %0, %notlhs
+  br i1 %1, label %if.end3, label %if.then2
+
+if.then2:                                         ; preds = %entry
+  tail call void @foo(i64 %A, i64 %B)
+  br label %if.end3
+
+if.end3:                                          ; preds = %if.then2, %entry
+  ret void
+}
+
+; CHECK-LABEL: test2
+; CHECK: cbz {{x[0-9]}}, {{.LBB1_3}}
+; CHECK: tbz [[REG1:x[0-9]+]], #3, {{.LBB1_3}}
+; CHECK-NOT: and [REG2:x[0-9]+], [[REG1]], #0x08
+; CHECK-NEXT-NOT: cbz [[REG2]], {{.LBB1_3}}
+
+define void @test2(i64 %A, i64* readonly %B) #0 {
+entry:
+  %tobool = icmp eq i64* %B, null
+  %and = and i64 %A, 8
+  %tobool1 = icmp eq i64 %and, 0
+  %or.cond = or i1 %tobool, %tobool1
+  br i1 %or.cond, label %if.end3, label %if.then2
+
+if.then2:                                         ; preds = %entry
+  %0 = load i64, i64* %B, align 4
+  tail call void @foo(i64 %A, i64 %0)
+  br label %if.end3
+
+if.end3:                                          ; preds = %entry, %if.then2
+  ret void
+}
+
+
+declare void @foo(i64, i64)

test/CodeGen/AArch64/fast-isel-tbz.ll

-; RUN: llc                             -aarch64-atomic-cfg-tidy=0 -verify-machineinstrs -mtriple=aarch64-apple-darwin < %s | FileCheck --check-prefix=CHECK %s
+; RUN: llc -disable-peephole                            -aarch64-atomic-cfg-tidy=0 -verify-machineinstrs -mtriple=aarch64-apple-darwin < %s | FileCheck --check-prefix=CHECK %s
 ; RUN: llc -fast-isel -fast-isel-abort=1 -aarch64-atomic-cfg-tidy=0 -verify-machineinstrs -mtriple=aarch64-apple-darwin < %s | FileCheck --check-prefix=CHECK --check-prefix=FAST %s
 
 define i32 @icmp_eq_i8(i8 zeroext %a) {
 ; CHECK-LABEL: icmp_eq_i8
 ; CHECK:       tbz {{w[0-9]+}}, #0, {{LBB.+_2}}
   %1 = and i8 %a, 1
   %2 = icmp eq i8 %1, 0
   br i1 %2, label %bb1, label %bb2, !prof !0