[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

2980

s/test/compare/

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

3039

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.

3053

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

3076

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?

3078

Single statement; No need for curly brackets.

You could also use a select statement.

3085

Same.

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.cpp

3038–3039

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.

3040

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?

[bmakam]

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

lib/Target/AArch64/AArch64InstrInfo.cpp

3038–3039

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

3040

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

3076

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.

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.cpp

3038–3039

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 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 {
   MachineBasicBlock *MBB = MI->getParent();
   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 *>>

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