[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 ↗

(On Diff #50038)

s/test/compare/

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

3015 ↗

(On Diff #50038)

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 ↗

(On Diff #50038)

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

3052 ↗

(On Diff #50038)

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 ↗

(On Diff #50038)

Single statement; No need for curly brackets.

You could also use a select statement.

3061 ↗

(On Diff #50038)

Same.

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.cpp

3016 ↗

(On Diff #50038)

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 ↗

(On Diff #50038)

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 ↗

(On Diff #50038)

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

3052 ↗

(On Diff #50038)

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 ↗

(On Diff #50038)

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 ↗

(On Diff #50038)

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

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

[mcrosier]

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

2980

Committed r263942.

llvm/trunk/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 compare and branch sequence by TBZ/TBNZ instruction when
+/// the compare'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 through COPY instructions to find definition.
+  while (DefMI->isCopy()) {
+    unsigned CopyVReg = DefMI->getOperand(1).getReg();
+    if (!MRI->hasOneNonDBGUse(CopyVReg))
+      return false;
+    if (!MRI->hasOneDef(CopyVReg))
+      return false;
+    DefMI = MRI->getVRegDef(CopyVReg);
+  }
+
+  switch (DefMI->getOpcode()) {
+  default:
+    return false;
+  // Fold AND into a TBZ/TBNZ if constant operand is power of 2.
+  case AArch64::ANDWri:
+  case AArch64::ANDXri: {
+    if (IsTestAndBranch)
+      return false;
+    if (DefMI->getParent() != MBB)
+      return false;
+    if (!MRI->hasOneNonDBGUse(VReg))
+      return false;
+
+    uint64_t Mask = AArch64_AM::decodeLogicalImmediate(
+        DefMI->getOperand(2).getImm(),
+        (DefMI->getOpcode() == AArch64::ANDWri) ? 32 : 64);
+    if (!isPowerOf2_64(Mask))
+      return false;
+
+    MachineOperand &MO = DefMI->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 = (Imm < 32)
+                       ? (IsNegativeBranch ? AArch64::TBNZW : AArch64::TBZW)
+                       : (IsNegativeBranch ? AArch64::TBNZX : AArch64::TBZX);
+    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 &&
+  case AArch64::CSINCWr:
+  case AArch64::CSINCXr: {
+    if (!(DefMI->getOperand(1).getReg() == AArch64::WZR &&
           DefMI->getOperand(2).getReg() == AArch64::WZR) &&
-      !(DefMI->getOpcode() == AArch64::CSINCXr &&
-        DefMI->getOperand(1).getReg() == AArch64::XZR &&
+        !(DefMI->getOperand(1).getReg() == AArch64::XZR &&
           DefMI->getOperand(2).getReg() == AArch64::XZR))
       return false;
 
     if (DefMI->findRegisterDefOperandIdx(AArch64::NZCV, true) != -1)
       return false;
 
-  AArch64CC::CondCode CC =
-      (AArch64CC::CondCode)DefMI->getOperand(3).getImm();
+    AArch64CC::CondCode CC = (AArch64CC::CondCode)DefMI->getOperand(3).getImm();
     bool CheckOnlyCCWrites = true;
     // Convert only when the condition code is not modified between
     // the CSINC and the branch. The CC may be used by other
     // instructions in between.
     if (modifiesConditionCode(DefMI, MI, CheckOnlyCCWrites, &getRegisterInfo()))
       return false;
     MachineBasicBlock &RefToMBB = *MBB;
     MachineBasicBlock *TBB = MI->getOperand(TargetBBInMI).getMBB();
     DebugLoc DL = MI->getDebugLoc();
     if (IsNegativeBranch)
       CC = AArch64CC::getInvertedCondCode(CC);
     BuildMI(RefToMBB, MI, DL, get(AArch64::Bcc)).addImm(CC).addMBB(TBB);
     MI->eraseFromParent();
     return true;
   }
+  }
+}
 
 std::pair<unsigned, unsigned>
 AArch64InstrInfo::decomposeMachineOperandsTargetFlags(unsigned TF) const {
   const unsigned Mask = AArch64II::MO_FRAGMENT;
   return std::make_pair(TF & Mask, TF & ~Mask);
 }
 
 ArrayRef<std::pair<unsigned, const char *>>

llvm/trunk/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)

llvm/trunk/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