[GlobalISel][AArch64] Select CBZ.

CBZ/CBNZ represent a substantial portion of all conditional branches.
Look through G_ICMP to select them.

We can't use tablegen yet because the existing patterns match an
AArch64ISD node.

llvm-svn: 298856

3 files changed, 161 insertions, 3 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
index 6490d88a3c8..2e199d9c0d9 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -448,6 +448,50 @@ static void changeFCMPPredToAArch64CC(CmpInst::Predicate P,
   }
 }
 
+bool AArch64InstructionSelector::selectCompareBranch(
+    MachineInstr &I, MachineFunction &MF, MachineRegisterInfo &MRI) const {
+
+  const unsigned CondReg = I.getOperand(0).getReg();
+  MachineBasicBlock *DestMBB = I.getOperand(1).getMBB();
+  MachineInstr *CCMI = MRI.getVRegDef(CondReg);
+  if (CCMI->getOpcode() != TargetOpcode::G_ICMP)
+    return false;
+
+  unsigned LHS = CCMI->getOperand(2).getReg();
+  unsigned RHS = CCMI->getOperand(3).getReg();
+  if (!getConstantVRegVal(RHS, MRI))
+    std::swap(RHS, LHS);
+
+  const auto RHSImm = getConstantVRegVal(RHS, MRI);
+  if (!RHSImm || *RHSImm != 0)
+    return false;
+
+  const RegisterBank &RB = *RBI.getRegBank(LHS, MRI, TRI);
+  if (RB.getID() != AArch64::GPRRegBankID)
+    return false;
+
+  const auto Pred = (CmpInst::Predicate)CCMI->getOperand(1).getPredicate();
+  if (Pred != CmpInst::ICMP_NE && Pred != CmpInst::ICMP_EQ)
+    return false;
+
+  const unsigned CmpWidth = MRI.getType(LHS).getSizeInBits();
+  unsigned CBOpc = 0;
+  if (CmpWidth <= 32)
+    CBOpc = (Pred == CmpInst::ICMP_EQ ? AArch64::CBZW : AArch64::CBNZW);
+  else if (CmpWidth == 64)
+    CBOpc = (Pred == CmpInst::ICMP_EQ ? AArch64::CBZX : AArch64::CBNZX);
+  else
+    return false;
+
+  auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(CBOpc))
+                 .addUse(LHS)
+                 .addMBB(DestMBB);
+
+  constrainSelectedInstRegOperands(*MIB.getInstr(), TII, TRI, RBI);
+  I.eraseFromParent();
+  return true;
+}
+
 bool AArch64InstructionSelector::selectVaStartAAPCS(
     MachineInstr &I, MachineFunction &MF, MachineRegisterInfo &MRI) const {
   return false;
@@ -556,6 +600,9 @@ bool AArch64InstructionSelector::select(MachineInstr &I) const {
     const unsigned CondReg = I.getOperand(0).getReg();
     MachineBasicBlock *DestMBB = I.getOperand(1).getMBB();
 
+    if (selectCompareBranch(I, MF, MRI))
+      return true;
+
     auto MIB = BuildMI(MBB, I, I.getDebugLoc(), TII.get(AArch64::TBNZW))
                    .addUse(CondReg)
                    .addImm(/*bit offset=*/0)
diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.h b/llvm/lib/Target/AArch64/AArch64InstructionSelector.h
index f8472f7a06d..c5476c4fbb7 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.h
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.h
@@ -40,14 +40,17 @@ public:
   bool select(MachineInstr &I) const override;
 
 private:
+  /// tblgen-erated 'select' implementation, used as the initial selector for
+  /// the patterns that don't require complex C++.
+  bool selectImpl(MachineInstr &I) const;
+
   bool selectVaStartAAPCS(MachineInstr &I, MachineFunction &MF,
                           MachineRegisterInfo &MRI) const;
   bool selectVaStartDarwin(MachineInstr &I, MachineFunction &MF,
                            MachineRegisterInfo &MRI) const;
 
-  /// tblgen-erated 'select' implementation, used as the initial selector for
-  /// the patterns that don't require complex C++.
-  bool selectImpl(MachineInstr &I) const;
+  bool selectCompareBranch(MachineInstr &I, MachineFunction &MF,
+                           MachineRegisterInfo &MRI) const;
 
   bool selectArithImmed(MachineOperand &Root, MachineOperand &Result1,
                         MachineOperand &Result2) const;
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-cbz.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-cbz.mir
new file mode 100644
index 00000000000..2decb994b96
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-cbz.mir
@@ -0,0 +1,108 @@
+# RUN: llc -mtriple=aarch64-- -run-pass=instruction-select -verify-machineinstrs -global-isel %s -o - | FileCheck %s
+
+--- |
+  define void @cbz_s32() { ret void }
+  define void @cbz_s64() { ret void }
+  define void @cbnz_s32() { ret void }
+  define void @cbnz_s64() { ret void }
+...
+
+---
+# CHECK-LABEL: name: cbz_s32
+name:            cbz_s32
+legalized:       true
+regBankSelected: true
+
+# CHECK:  body:
+# CHECK:   bb.0:
+# CHECK:    %0 = COPY %w0
+# CHECK:    CBZW %0, %bb.1
+# CHECK:    B %bb.0
+body:             |
+  bb.0:
+    liveins: %w0
+    successors: %bb.0, %bb.1
+
+    %0:gpr(s32) = COPY %w0
+    %1:gpr(s32) = G_CONSTANT i32 0
+    %2:gpr(s1) = G_ICMP intpred(eq), %0, %1
+    G_BRCOND %2(s1), %bb.1
+    G_BR %bb.0
+
+  bb.1:
+...
+
+---
+# CHECK-LABEL: name: cbz_s64
+name:            cbz_s64
+legalized:       true
+regBankSelected: true
+
+# CHECK:  body:
+# CHECK:   bb.0:
+# CHECK:    %0 = COPY %x0
+# CHECK:    CBZX %0, %bb.1
+# CHECK:    B %bb.0
+body:             |
+  bb.0:
+    liveins: %x0
+    successors: %bb.0, %bb.1
+
+    %0:gpr(s64) = COPY %x0
+    %1:gpr(s64) = G_CONSTANT i64 0
+    %2:gpr(s1) = G_ICMP intpred(eq), %0, %1
+    G_BRCOND %2(s1), %bb.1
+    G_BR %bb.0
+
+  bb.1:
+...
+
+---
+# CHECK-LABEL: name: cbnz_s32
+name:            cbnz_s32
+legalized:       true
+regBankSelected: true
+
+# CHECK:  body:
+# CHECK:   bb.0:
+# CHECK:    %0 = COPY %w0
+# CHECK:    CBNZW %0, %bb.1
+# CHECK:    B %bb.0
+body:             |
+  bb.0:
+    liveins: %w0
+    successors: %bb.0, %bb.1
+
+    %0:gpr(s32) = COPY %w0
+    %1:gpr(s32) = G_CONSTANT i32 0
+    %2:gpr(s1) = G_ICMP intpred(ne), %0, %1
+    G_BRCOND %2(s1), %bb.1
+    G_BR %bb.0
+
+  bb.1:
+...
+
+---
+# CHECK-LABEL: name: cbnz_s64
+name:            cbnz_s64
+legalized:       true
+regBankSelected: true
+
+# CHECK:  body:
+# CHECK:   bb.0:
+# CHECK:    %0 = COPY %x0
+# CHECK:    CBNZX %0, %bb.1
+# CHECK:    B %bb.0
+body:             |
+  bb.0:
+    liveins: %x0
+    successors: %bb.0, %bb.1
+
+    %0:gpr(s64) = COPY %x0
+    %1:gpr(s64) = G_CONSTANT i64 0
+    %2:gpr(s1) = G_ICMP intpred(ne), %0, %1
+    G_BRCOND %2(s1), %bb.1
+    G_BR %bb.0
+
+  bb.1:
+...