[AArch64][CodeGen] Peephole optimization of Cmp+Bcc sequences with unconditional branch

[eastig]

Summary

This is a peephole optimization of Cmp+Bcc sequences with an unconditional branch. Such sequences can appear after the patch: http://reviews.llvm.org/D18838

When we compare with zero a result can be known. We can figure out which path is taken by Bcc and substitute Bcc by an unconditional branch to that path.

Rules:

SUBS reg, 0; B.LO              => B <false path>
SUBS reg, 0; B.HS <label> => B <label>

[jmolloy]

Hi Evgeny,

I'm afraid I don't understand the transform from the description. This looks like a simple compare-with-zero pattern - how do we know which branch will be taken?

James

[eastig]

In D22226#480417, @jmolloy wrote:

Hi Evgeny,

I'm afraid I don't understand the transform from the description. This looks like a simple compare-with-zero pattern - how do we know which branch will be taken?

James

'SUBS WZR, Wn, #0' always sets C to 1 because nothing is borrowed. So we can say where B.LO and B.HS will go: B.LO will fall through (C is not clear); B.HS will take its target (C is set).
IR, which corresponds this, looks like:

%2 = icmp ult/uge i64 %1, 0
 br i1 %2, label %ret1, label %ret0

BTW it is not optimized in the middle-end. In the back-end the optimized sequence 'SUBS; Bcc' is a result of some other optimizations. In a test, we have seen this case, there is a sequence of pointer expressions and a comparison.

lib/Target/AArch64/AArch64InstrInfo.h

public:
   getSerializableBitmaskMachineOperandTargetFlags() const override;
 
 private:
   void instantiateCondBranch(MachineBasicBlock &MBB, const DebugLoc &DL,
                              MachineBasicBlock *TBB,
                              ArrayRef<MachineOperand> Cond) const;
   bool substituteCmpToZero(MachineInstr &CmpInstr, unsigned SrcReg,
                            const MachineRegisterInfo *MRI) const;
+  bool optimizeCmpAndBccSeq(MachineInstr &BrcInstr) const;
 };
 
 /// emitFrameOffset - Emit instructions as needed to set DestReg to SrcReg
 /// plus Offset.  This is intended to be used from within the prolog/epilog
 /// insertion (PEI) pass, where a virtual scratch register may be allocated
 /// if necessary, to be replaced by the scavenger at the end of PEI.
 void emitFrameOffset(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                      const DebugLoc &DL, unsigned DestReg, unsigned SrcReg,

lib/Target/AArch64/AArch64InstrInfo.cpp

 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/PseudoSourceValue.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/TargetRegistry.h"
 #include <algorithm>
+#include <tuple>
 
 using namespace llvm;
 
 #define GET_INSTRINFO_CTOR_DTOR
 #include "AArch64GenInstrInfo.inc"
 
 AArch64InstrInfo::AArch64InstrInfo(const AArch64Subtarget &STI)
     : AArch64GenInstrInfo(AArch64::ADJCALLSTACKDOWN, AArch64::ADJCALLSTACKUP),
void AArch64InstrInfo::genAlternativeCodeSequence(
   } // end switch (Pattern)
   // Record MUL and ADD/SUB for deletion
   DelInstrs.push_back(MUL);
   DelInstrs.push_back(&Root);
 
   return;
 }
 
+static bool isCmpToZero(MachineInstr &CmpInstr) {
+  return isSUBSRegImm(CmpInstr.getOpcode())
+         && (CmpInstr.getOperand(2).getImm() == 0);
+}
+
+static MachineRegisterInfo *getMRI(MachineInstr &MI) {
+  assert(MI.getParent() && "Incomplete machine instruciton\n");
+  assert(MI.getParent()->getParent());
+  return &MI.getParent()->getParent()->getRegInfo();
+}
+
+/// Returns basic blocks which correspond to True- and False paths of the
+/// specified conditional branch.
+static std::tuple<MachineBasicBlock*, MachineBasicBlock*>
+findTrueAndFalseTargets(MachineInstr &BrcInstr) {
+  assert(BrcInstr.getOpcode() == AArch64::Bcc);
+  MachineBasicBlock *TBB = BrcInstr.getOperand(1).getMBB();
+  // Find fallthrough basic block
+  auto BBI = std::find_if(BrcInstr.getParent()->succ_begin(),
+                          BrcInstr.getParent()->succ_end(),
+                          [TBB](MachineBasicBlock *BB) {
+                            assert(BB);
+                            return !BB->isEHPad() && BB != TBB;
+                          });
+  assert(BBI != BrcInstr.getParent()->succ_end());
+  return std::make_tuple(TBB, *BBI);
+}
+
+/// Optimize Cmp and Bcc sequences.
+/// When we compare with zero a result is known. We can figure out
+/// which path is taken by Brc. We substitute Brc by an unconditional
+/// branch to that path.
+///
+/// Rules:
+///  SUBS reg, 0; B.LO         => B <false path>
+///  SUBS reg, 0; B.HS <label> => B <label>
+bool AArch64InstrInfo::optimizeCmpAndBccSeq(MachineInstr &BrcInstr) const {
+  const TargetRegisterInfo *TRI = &getRegisterInfo();
+  auto CmpInstrLoc = std::find_if(MachineBasicBlock::reverse_iterator(&BrcInstr),
+          BrcInstr.getParent()->rend(),
+          [TRI](MachineInstr &MI) {
+            return MI.modifiesRegister(AArch64::NZCV, TRI);
+          });
+
+  if (CmpInstrLoc == BrcInstr.getParent()->rend())
+    return false;
+
+  MachineInstr &CmpInstr = *CmpInstrLoc;
+  assert(getMRI(CmpInstr));
+  if (!isCmpToZero(CmpInstr)
+      || !getMRI(CmpInstr)->use_nodbg_empty(CmpInstr.getOperand(0).getReg())
+      || areCFlagsAccessedBetweenInstrs(&CmpInstr, &BrcInstr, TRI, AK_Read)
+      || areCFlagsAliveInSuccessors(BrcInstr.getParent()))
+    return false;
+
+  MachineBasicBlock *UncondBrTarget = nullptr;
+  MachineBasicBlock *TargetToRemove = nullptr;
+  switch (static_cast<AArch64CC::CondCode>(BrcInstr.getOperand(0).getImm())) {
+    case AArch64CC::HS:
+      // SUBS reg, 0; B.HS <true path> => B <true path>
+      std::tie(UncondBrTarget, TargetToRemove) = findTrueAndFalseTargets(BrcInstr);
+      break;
+      
+    case AArch64CC::LO:
+      // SUBS reg, 0; B.LO => B <false path>
+      std::tie(TargetToRemove, UncondBrTarget) = findTrueAndFalseTargets(BrcInstr);
+      break;
+
+    default:
+      return false;
+  }
+
+  assert(UncondBrTarget);
+  assert(TargetToRemove);
+
+  BrcInstr.getParent()->removeSuccessor(TargetToRemove);
+  while (BrcInstr.getNumOperands())
+    BrcInstr.RemoveOperand(BrcInstr.getNumOperands() - 1);
+  BrcInstr.setDesc(get(AArch64::B));
+  BrcInstr.addOperand(MachineOperand::CreateMBB(UncondBrTarget));
+  CmpInstr.eraseFromParent();
+  return true;
+}
+
 /// \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>

 bool AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) const {
   bool IsNegativeBranch = false;
   bool IsTestAndBranch = false;
   unsigned TargetBBInMI = 0;
   switch (MI.getOpcode()) {
   default:
     llvm_unreachable("Unknown branch instruction?");
   case AArch64::Bcc:
-    return false;
+    return optimizeCmpAndBccSeq(MI);
   case AArch64::CBZW:
   case AArch64::CBZX:
     TargetBBInMI = 1;
     break;
   case AArch64::CBNZW:
   case AArch64::CBNZX:
     TargetBBInMI = 1;
     IsNegativeBranch = true;

test/CodeGen/AArch64/cmp-bcc-opt.mir

+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass peephole-opt -o - %s 2>&1 | FileCheck --check-prefix CHECK-SUBS32-B_LO %s
+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass peephole-opt -o - %s 2>&1 | FileCheck --check-prefix CHECK-SUBS64-B_LO %s
+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass peephole-opt -o - %s 2>&1 | FileCheck --check-prefix CHECK-SUBS32-B_HS %s
+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass peephole-opt -o - %s 2>&1 | FileCheck --check-prefix CHECK-SUBS64-B_HS %s
+
+# CHECK-SUBS32-B_LO: %1 = ADDWri killed %1, 1, 0
+# CHECK-SUBS32-B_LO-NEXT: B %bb.2.ret0
+
+# CHECK-SUBS64-B_LO: %1 = ADDXri killed %1, 1, 0
+# CHECK-SUBS64-B_LO-NEXT: B %bb.2.ret0
+
+# CHECK-SUBS32-B_HS: %1 = ADDWri killed %1, 2, 0
+# CHECK-SUBS32-B_HS-NEXT: B %bb.1.ret1
+
+# CHECK-SUBS64-B_HS: %1 = ADDXri killed %1, 2, 0
+# CHECK-SUBS64-B_HS-NEXT: B %bb.1.ret1
+
+--- |
+  define i32 @test01() nounwind {
+  entry:
+    %0 = select i1 true, i32 1, i32 0
+    %1 = add i32 %0, 1
+    %2 = icmp ult i32 %1, 0
+    br i1 %2, label %ret1, label %ret0
+
+  ret1:                                      ; preds = %entry
+    ret i32 1
+
+  ret0:                                      ; preds = %entry
+    ret i32 0
+  }
+
+  define i32 @test02() nounwind {
+  entry:
+    %0 = select i1 true, i64 1, i64 0
+    %1 = add i64 %0, 1
+    %2 = icmp ult i64 %1, 0
+    br i1 %2, label %ret1, label %ret0
+
+  ret1:                                      ; preds = %entry
+    ret i32 1
+
+  ret0:                                      ; preds = %entry
+    ret i32 0
+  }
+
+  define i32 @test03() nounwind {
+  entry:
+    %0 = select i1 true, i32 1, i32 0
+    %1 = add i32 %0, 2
+    %2 = icmp uge i32 %1, 0
+    br i1 %2, label %ret1, label %ret0
+
+  ret1:                                      ; preds = %entry
+    ret i32 1
+
+  ret0:                                      ; preds = %entry
+    ret i32 0
+  }
+
+  define i32 @test04() nounwind {
+  entry:
+    %0 = select i1 true, i64 1, i64 0
+    %1 = add i64 %0, 2
+    %2 = icmp uge i64 %1, 0
+    br i1 %2, label %ret1, label %ret0
+
+  ret1:                                      ; preds = %entry
+    ret i32 1
+
+  ret0:                                      ; preds = %entry
+    ret i32 0
+  }
+...
+---
+name:            test01
+registers:
+  - { id: 0, class: gpr32 }
+  - { id: 1, class: gpr32common }
+body:             |
+  bb.0.entry:
+    successors: %bb.2.ret0, %bb.1.ret1
+
+    %0 = MOVi32imm 1
+    %1 = ADDWri killed %1, 1, 0
+    %wzr = SUBSWri killed %1, 0, 0, implicit-def %nzcv
+    Bcc 3, %bb.1.ret1, implicit %nzcv
+    B %bb.2.ret0
+
+  bb.1.ret1:
+    %w0 = MOVi32imm 1
+    RET_ReallyLR implicit %w0
+
+  bb.2.ret0:
+    %w0 = MOVi32imm 0
+    RET_ReallyLR implicit %w0
+...
+---
+name:            test02
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64common }
+body:             |
+  bb.0.entry:
+    successors: %bb.2.ret0, %bb.1.ret1
+
+    %0 = MOVi64imm 1
+    %1 = ADDXri killed %1, 1, 0
+    %xzr = SUBSXri killed %1, 0, 0, implicit-def %nzcv
+    Bcc 3, %bb.1.ret1, implicit %nzcv
+    B %bb.2.ret0
+
+  bb.1.ret1:
+    %w0 = MOVi32imm 1
+    RET_ReallyLR implicit %w0
+
+  bb.2.ret0:
+    %w0 = MOVi32imm 0
+    RET_ReallyLR implicit %w0
+
+...
+---
+name:            test03
+registers:
+  - { id: 0, class: gpr32 }
+  - { id: 1, class: gpr32common }
+body:             |
+  bb.0.entry:
+    successors: %bb.2.ret0, %bb.1.ret1
+
+    %0 = MOVi32imm 1
+    %1 = ADDWri killed %1, 2, 0
+    %wzr = SUBSWri killed %1, 0, 0, implicit-def %nzcv
+    Bcc 2, %bb.1.ret1, implicit %nzcv
+    B %bb.2.ret0
+
+  bb.1.ret1:
+    %w0 = MOVi32imm 1
+    RET_ReallyLR implicit %w0
+
+  bb.2.ret0:
+    %w0 = MOVi32imm 0
+    RET_ReallyLR implicit %w0
+...
+---
+name:            test04
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64common }
+body:             |
+  bb.0.entry:
+    successors: %bb.2.ret0, %bb.1.ret1
+
+    %0 = MOVi64imm 1
+    %1 = ADDXri killed %1, 2, 0
+    %xzr = SUBSXri killed %1, 0, 0, implicit-def %nzcv
+    Bcc 2, %bb.1.ret1, implicit %nzcv
+    B %bb.2.ret0
+
+  bb.1.ret1:
+    %w0 = MOVi32imm 1
+    RET_ReallyLR implicit %w0
+
+  bb.2.ret0:
+    %w0 = MOVi32imm 0
+    RET_ReallyLR implicit %w0
+
+...
+