[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

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..?

3006

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.

3020

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

3043

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?

3045

Single statement; No need for curly brackets.

You could also use a select statement.

3052

Same.

[t.p.northover]

lib/Target/AArch64/AArch64InstrInfo.cpp

3007

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?

3007

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

3007

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

3007

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

3043

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

3007

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.

include/llvm/Target/TargetInstrInfo.h

public:
   virtual bool optimizeCompareInstr(MachineInstr *CmpInstr,
                                     unsigned SrcReg, unsigned SrcReg2,
                                     int Mask, int Value,
                                     const MachineRegisterInfo *MRI) const {
     return false;
   }
   virtual bool optimizeCondBranch(MachineInstr *MI) const { return false; }
 
+  virtual bool optimizeTestAndBranch(MachineInstr *MI) const {return false; }
+
   /// Try to remove the load by folding it to a register operand at the use.
   /// We fold the load instructions if and only if the
   /// def and use are in the same BB. We only look at one load and see
   /// whether it can be folded into MI. FoldAsLoadDefReg is the virtual register
   /// defined by the load we are trying to fold. DefMI returns the machine
   /// instruction that defines FoldAsLoadDefReg, and the function returns
   /// the machine instruction generated due to folding.
   virtual MachineInstr* optimizeLoadInstr(MachineInstr *MI,

lib/CodeGen/PeepholeOptimizer.cpp

namespace {
 
   private:
     bool optimizeCmpInstr(MachineInstr *MI, MachineBasicBlock *MBB);
     bool optimizeExtInstr(MachineInstr *MI, MachineBasicBlock *MBB,
                           SmallPtrSetImpl<MachineInstr*> &LocalMIs);
     bool optimizeSelect(MachineInstr *MI,
                         SmallPtrSetImpl<MachineInstr *> &LocalMIs);
     bool optimizeCondBranch(MachineInstr *MI);
+    bool optimizeTestAndBranch(MachineInstr *MI);
     bool optimizeCoalescableCopy(MachineInstr *MI);
     bool optimizeUncoalescableCopy(MachineInstr *MI,
                                    SmallPtrSetImpl<MachineInstr *> &LocalMIs);
     bool findNextSource(unsigned Reg, unsigned SubReg,
                         RewriteMapTy &RewriteMap);
     bool isMoveImmediate(MachineInstr *MI,
                          SmallSet<unsigned, 4> &ImmDefRegs,
                          DenseMap<unsigned, MachineInstr*> &ImmDefMIs);

 }
 
 /// \brief Check if a simpler conditional branch can be
 // generated
 bool PeepholeOptimizer::optimizeCondBranch(MachineInstr *MI) {
   return TII->optimizeCondBranch(MI);
 }
 
+/// \brief Check if a TBZ/TBNZ can be
+// generated instead of an AND followed by CBZ/CBNZ.
+bool PeepholeOptimizer::optimizeTestAndBranch(MachineInstr *MI) {
+  return TII->optimizeTestAndBranch(MI);
+}
+
 /// \brief Try to find the next source that share the same register file
 /// for the value defined by \p Reg and \p SubReg.
 /// When true is returned, the \p RewriteMap can be used by the client to
 /// retrieve all Def -> Use along the way up to the next source. Any found
 /// Use that is not itself a key for another entry, is the next source to
 /// use. During the search for the next source, multiple sources can be found
 /// given multiple incoming sources of a PHI instruction. In this case, we
 /// look in each PHI source for the next source; all found next sources must
for (MachineBasicBlock::iterator MII = MBB.begin(), MIE = MBB.end();
           (MI->isCompare() && optimizeCmpInstr(MI, &MBB)) ||
           (MI->isSelect() && optimizeSelect(MI, LocalMIs))) {
         // MI is deleted.
         LocalMIs.erase(MI);
         Changed = true;
         continue;
       }
 
+      if (MI->isConditionalBranch() && optimizeTestAndBranch(MI)) {
+        Changed = true;
+        continue;
+      }
+
       if (MI->isConditionalBranch() && optimizeCondBranch(MI)) {
         Changed = true;
         continue;
       }
 
       if (isCoalescableCopy(*MI) && optimizeCoalescableCopy(MI)) {
         // MI is just rewritten.
         Changed = true;

lib/Target/AArch64/AArch64InstrInfo.h

bool analyzeCompare(const MachineInstr *MI, unsigned &SrcReg,
                       unsigned &SrcReg2, int &CmpMask,
                       int &CmpValue) const override;
   /// optimizeCompareInstr - Convert the instruction supplying the argument to
   /// the comparison into one that sets the zero bit in the flags register.
   bool optimizeCompareInstr(MachineInstr *CmpInstr, unsigned SrcReg,
                             unsigned SrcReg2, int CmpMask, int CmpValue,
                             const MachineRegisterInfo *MRI) const override;
   bool optimizeCondBranch(MachineInstr *MI) const override;
+  bool optimizeTestAndBranch(MachineInstr *MI) const override;
   /// Return true when there is potentially a faster code sequence
   /// for an instruction chain ending in <Root>. All potential patterns are
   /// listed in the <Patterns> array.
   bool getMachineCombinerPatterns(MachineInstr &Root,
                   SmallVectorImpl<MachineCombinerPattern> &Patterns)
       const override;
   /// Return true when Inst is associative and commutative so that it can be
   /// reassociated.

lib/Target/AArch64/AArch64InstrInfo.cpp

 ///   b.<inverted condition code>
 ///
 /// 2.
 ///   csinc w9, wzr, wzr, <condition code>
 ///   tbz   w9, #0, 0x44
 /// to
 ///   b.<condition code>
 ///
 /// \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()) {
   default:
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 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)
     return false;
 
   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;
 }
 
+/// \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 TBZ/TBNZ is generated
+///
+bool AArch64InstrInfo::optimizeTestAndBranch(MachineInstr *MI) const {
+  bool Is64Bit = false;
+  bool IsNegativeBranch = false;
+  switch (MI->getOpcode()) {
+  default:
+    return false;
+  case AArch64::CBZW:
+    break;
+  case AArch64::CBNZW:
+    IsNegativeBranch = true;
+    break;
+  case AArch64::CBZX:
+    Is64Bit = true;
+    break;
+  case AArch64::CBNZX:
+    IsNegativeBranch = true;
+    Is64Bit = true;
+    break;
+  }
+  // Find Definition.
+  assert(MI->getParent() && "Incomplete machine instruciton\n");
+  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 for AND
+  if (DefMI->getParent() != MBB) return false;
+
+  if (!DefMI->isCopy()) 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(),
+                  Is64Bit ? 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;
+}
+
 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 *>>
 AArch64InstrInfo::getSerializableDirectMachineOperandTargetFlags() const {

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