ARM64 - Conditional branch peephole

[Gerolf]

Summary

Converts CSINC-branch sequence to single conditional branch

On ARM64 the optimization generates a single conditional branch
for csinc-branch sequences like in the examples below. This is
possible when the csinc sets or clears a register based on a condition
code and the branch checks that register. Also the condition
code may not be modified between the CINCR and the original branch.

Examples:

1. csinc w9, wzr, wzr, <condition code> tbnz w9, #0, 0x44 to b.<inverted condition code>
   
   2. csinc w9, wzr, wzr, <condition code> tbz w9, #0, 0x44 to b.<condition code>

[rengolin]

Hi Gerolf,

Overall, the patch looks good. I think the substitution is sound, but I'd rather Tim had a look to confirm.

There are a few in-line comments that need addressing...

cheers,
--renato

lib/CodeGen/PeepholeOptimizer.cpp

506

Couldn't you just...

return TII->optimizeCondBranch(MI);

?

lib/Target/AArch64/AArch64InstrInfo.cpp

2864

Same here with comments. If you call your variables more explicitly, you don't need the one-line comment before.

Ideas are:

bool IsNegativeBranch;
bool IsTestAndBranch;
bool TargetsBBInMI;

Also, the coding standard mandates camel case for variable and function names, with caps on first letter for variables and low-caps for function names.

2868

Personal opinion, I don't like switches for this kind of thing. I think repeated ifs are more explicit to not repeat the assignments, at least it's more clear what each flag means from instructions. But I'm ok with the switch.

2877

unnecessary space

2899

I don't think it's possible for this to happen, since they're local variables and you're setting them above with an llvm_unreachable. These asserts are completely redundant.

2910

while I generally support asserts, I think you got too many of them here... It's making the rest of the code unreadable

2947

Did DefMI get inserted anywhere? I may be wrong, but I don't think this is necessary

[t.p.northover]

Hi Gerolf,

This looks reasonable to me, with a couple of tweaks. I agree with most of Reanto's points too.

Cheers.

Tim.

lib/Target/AArch64/AArch64InstrInfo.cpp

776

Transplanted comment is out of date.

2939–2940

We usually prefer to bail early in failure cases:

if (modifiesConditionCode(...))
  return false;
[...]

2947

It's already going to be in the basic block, isn't it? It's not one we created.

[Gerolf]

Renato and Tim thanks for your feedback. I followed up on all but the switch vs if comment.

lib/CodeGen/PeepholeOptimizer.cpp

506

done :-)

lib/Target/AArch64/AArch64InstrInfo.cpp

2864

yes, thanks.

2868

Thanks. I'm not in favor of such code either, but don't see significant improvement with the ifs. Laziness wins! :-)

2877

removed.

2899

removed.

2910

agree. Removed all but one.

2939–2940

The only reason I did this way was was to make Build happy that need a reference to MBB. But I can introduce a new variable for it and follow the preference.

2947

Yes, that line should go or be fixed: It should be dead code when MI is removed. If it is not dead code it needs an additional check that MI is the only use.

[Gerolf]

Update based on feedback from Renato and Tim.

[rengolin]

Thanks Gerolf, LGTM.

include/llvm/Target/TargetInstrInfo.h

public:
   /// into something more efficient. E.g., on ARM most instructions can set the
   /// flags register, obviating the need for a separate CMP.
   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; }
 
   /// optimizeLoadInstr - 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.

lib/CodeGen/PeepholeOptimizer.cpp

void getAnalysisUsage(AnalysisUsage &AU) const override {
       }
     }
 
   private:
     bool optimizeCmpInstr(MachineInstr *MI, MachineBasicBlock *MBB);
     bool optimizeExtInstr(MachineInstr *MI, MachineBasicBlock *MBB,
                           SmallPtrSetImpl<MachineInstr*> &LocalMIs);
     bool optimizeSelect(MachineInstr *MI);
+    bool optimizeCondBranch(MachineInstr *MI);
     bool optimizeCopyOrBitcast(MachineInstr *MI);
     bool optimizeCoalescableCopy(MachineInstr *MI);
     bool optimizeUncoalescableCopy(MachineInstr *MI,
                                    SmallPtrSetImpl<MachineInstr *> &LocalMIs);
     bool findNextSource(unsigned &Reg, unsigned &SubReg);
     bool isMoveImmediate(MachineInstr *MI,
                          SmallSet<unsigned, 4> &ImmDefRegs,
                          DenseMap<unsigned, MachineInstr*> &ImmDefMIs);
if (!Optimizable)
     return false;
   if (!TII->optimizeSelect(MI))
     return false;
   MI->eraseFromParent();
   ++NumSelects;
   return true;
 }
 
+/// \brief Check if a simpler conditional branch can be
+// generated
+bool PeepholeOptimizer::optimizeCondBranch(MachineInstr *MI) {
+  return TII->optimizeCondBranch(MI);
+}
+
 /// \brief Check if the registers defined by the pair (RegisterClass, SubReg)
 /// share the same register file.
 static bool shareSameRegisterFile(const TargetRegisterInfo &TRI,
                                   const TargetRegisterClass *DefRC,
                                   unsigned DefSubReg,
                                   const TargetRegisterClass *SrcRC,
                                   unsigned SrcSubReg) {
   // Same register class.
for (MachineBasicBlock::iterator
           (MI->isCompare() && optimizeCmpInstr(MI, MBB)) ||
           (MI->isSelect() && optimizeSelect(MI))) {
         // MI is deleted.
         LocalMIs.erase(MI);
         Changed = true;
         continue;
       }
 
+      if (MI->isConditionalBranch() && optimizeCondBranch(MI)) {
+        Changed = true;
+        continue;
+      }
+
       if (isCoalescableCopy(*MI) && optimizeCoalescableCopy(MI)) {
         // MI is just rewritten.
         Changed = true;
         continue;
       }
 
       if (isMoveImmediate(MI, ImmDefRegs, ImmDefMIs)) {
         SeenMoveImm = true;

lib/Target/AArch64/AArch64InstrInfo.h

public:
   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;
   /// hasPattern - return true when there is potentially a faster code sequence
   /// for an instruction chain ending in <Root>. All potential patterns are
   /// listed
   /// in the <Pattern> array.
   bool hasPattern(MachineInstr &Root,
                   SmallVectorImpl<MachineCombinerPattern::MC_PATTERN> &Pattern)
       const override;
 

lib/Target/AArch64/AArch64InstrInfo.cpp

static unsigned convertFlagSettingOpcode(MachineInstr *MI) {
     case AArch64::SUBSXrr:      NewOpc = AArch64::SUBXrr; break;
     case AArch64::SUBSXri:      NewOpc = AArch64::SUBXri; break;
     case AArch64::SUBSXrs:      NewOpc = AArch64::SUBXrs; break;
     case AArch64::SUBSXrx:      NewOpc = AArch64::SUBXrx; break;
     }
     return NewOpc;
 }
 
+/// True when condition code could be modified on the instruction
+/// trace starting at from and ending at to.
+static bool modifiesConditionCode(MachineInstr *From, MachineInstr *To,
+                                  const bool CheckOnlyCCWrites,
+                                  const TargetRegisterInfo *TRI) {
+  // We iterate backward starting \p To until we hit \p From
+  MachineBasicBlock::iterator I = To, E = From, B = To->getParent()->begin();
+
+  // Early exit if To is at the beginning of the BB.
+  if (I == B)
+    return true;
+
+  // Check whether the definition of SrcReg is in the same basic block as
+  // Compare. If not, assume the condition code gets modified on some path.
+  if (To->getParent() != From->getParent())
+    return true;
+
+  // Check that NZCV isn't set on the trace.
+  for (--I; I != E; --I) {
+    const MachineInstr &Instr = *I;
+
+    if (Instr.modifiesRegister(AArch64::NZCV, TRI) ||
+        (!CheckOnlyCCWrites && Instr.readsRegister(AArch64::NZCV, TRI)))
+      // This instruction modifies or uses NZCV after the one we want to
+      // change.
+      return true;
+    if (I == B)
+      // We currently don't allow the instruction trace to cross basic
+      // block boundaries
+      return true;
+  }
+  return false;
+}
 /// optimizeCompareInstr - Convert the instruction supplying the argument to the
 /// comparison into one that sets the zero bit in the flags register.
 bool AArch64InstrInfo::optimizeCompareInstr(
     MachineInstr *CmpInstr, unsigned SrcReg, unsigned SrcReg2, int CmpMask,
     int CmpValue, const MachineRegisterInfo *MRI) const {
 
   // Replace SUBSWrr with SUBWrr if NZCV is not used.
   int Cmp_NZCV = CmpInstr->findRegisterDefOperandIdx(AArch64::NZCV, true);
bool AArch64InstrInfo::optimizeCompareInstr(
   if (!MRI->use_nodbg_empty(CmpInstr->getOperand(0).getReg()))
     return false;
 
   // Get the unique definition of SrcReg.
   MachineInstr *MI = MRI->getUniqueVRegDef(SrcReg);
   if (!MI)
     return false;
 
-  // We iterate backward, starting from the instruction before CmpInstr and
-  // stop when reaching the definition of the source register or done with the
-  // basic block, to check whether NZCV is used or modified in between.
-  MachineBasicBlock::iterator I = CmpInstr, E = MI,
-                              B = CmpInstr->getParent()->begin();
-
-  // Early exit if CmpInstr is at the beginning of the BB.
-  if (I == B)
-    return false;
-
-  // Check whether the definition of SrcReg is in the same basic block as
-  // Compare. If not, we can't optimize away the Compare.
-  if (MI->getParent() != CmpInstr->getParent())
-    return false;
-
-  // Check that NZCV isn't set between the comparison instruction and the one we
-  // want to change.
+  bool CheckOnlyCCWrites = false;
   const TargetRegisterInfo *TRI = &getRegisterInfo();
-  for (--I; I != E; --I) {
-    const MachineInstr &Instr = *I;
-
-    if (Instr.modifiesRegister(AArch64::NZCV, TRI) ||
-        Instr.readsRegister(AArch64::NZCV, TRI))
-      // This instruction modifies or uses NZCV after the one we want to
-      // change. We can't do this transformation.
-      return false;
-    if (I == B)
-      // The 'and' is below the comparison instruction.
-      return false;
-  }
+  if (modifiesConditionCode(MI, CmpInstr, CheckOnlyCCWrites, TRI))
 
   unsigned NewOpc = MI->getOpcode();
   switch (MI->getOpcode()) {
   default:
     return false;
   case AArch64::ADDSWrr:
   case AArch64::ADDSWri:
   case AArch64::ADDSXrr:
void AArch64InstrInfo::genAlternativeCodeSequence(
   }
   } // end switch (Pattern)
   // Record MUL and ADD/SUB for deletion
   DelInstrs.push_back(MUL);
   DelInstrs.push_back(&Root);
 
   return;
 }
+
+/// \brief Replace csincr-branch sequence by simple conditional branch
+///
+/// Examples:
+/// 1.
+///   csinc  w9, wzr, wzr, <condition code>
+///   tbnz   w9, #0, 0x44
+/// to
+///   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;
+  unsigned CCInMI = 0;
+  switch (MI->getOpcode()) {
+  default:
+    llvm_unreachable("Unknown branch instruction?");
+  case AArch64::Bcc:
+    return false;
+  case AArch64::CBZW:
+  case AArch64::CBZX:
+    TargetBBInMI = 1;
+    CCInMI = 2;
+    break;
+  case AArch64::CBNZW:
+  case AArch64::CBNZX:
+    TargetBBInMI = 1;
+    CCInMI = 2;
+    IsNegativeBranch = true;
+    break;
+  case AArch64::TBZW:
+  case AArch64::TBZX:
+    TargetBBInMI = 2;
+    CCInMI = 3;
+    IsTestAndBranch = true;
+    break;
+  case AArch64::TBNZW:
+  case AArch64::TBNZX:
+    TargetBBInMI = 2;
+    CCInMI = 3;
+    IsNegativeBranch = true;
+    IsTestAndBranch = true;
+    break;
+  }
+  // So we increment a zero register and test for bits other
+  // than bit 0? Conservatively bail out in case the verifier
+  // missed this case.
+  if (IsTestAndBranch && MI->getOperand(1).getImm())
+    return false;
+
+  // 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 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(CCInMI).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;
+}

test/CodeGen/AArch64/arm64-bcc.ll

+; RUN: llc < %s -mtriple=arm64-apple-darwint  | FileCheck %s
+; Checks for conditional branch b.vs
+
+; Function Attrs: nounwind
+define i32 @add(i32, i32) {
+entry:
+  %2 = tail call { i32, i1 } @llvm.sadd.with.overflow.i32(i32 %0, i32 %1)
+  %3 = extractvalue { i32, i1 } %2, 1
+  br i1 %3, label %6, label %4
+
+; <label>:4                                       ; preds = %entry
+  %5 = extractvalue { i32, i1 } %2, 0
+  ret i32 %5
+
+; <label>:6                                       ; preds = %entry
+  tail call void @llvm.trap()
+  unreachable
+; CHECK: b.vs
+}
+
+; Function Attrs: nounwind readnone
+declare { i32, i1 } @llvm.sadd.with.overflow.i32(i32, i32)
+
+; Function Attrs: noreturn nounwind
+declare void @llvm.trap()
+