[AArch64][GlobalISel][NFC] Factor out TB(N)Z emission code into its own function

[paquette]

Summary

Factor it out into emitTestBit and add some asserts to the new function.

This will be useful for implementing TB(N)Z emission for SLT/SGT compares.

[aemerson]

LGTM.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1135

type smaler

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

private:
   /// Otherwise, do nothing and return a nullptr.
   MachineInstr *emitFMovForFConstant(MachineInstr &MI,
                                      MachineRegisterInfo &MRI) const;
 
   /// Emit a CSet for a compare.
   MachineInstr *emitCSetForICMP(Register DefReg, unsigned Pred,
                                 MachineIRBuilder &MIRBuilder) const;
 
+  /// Emit a TB(N)Z instruction which tests \p Bit in \p TestReg.
+  /// \p IsNegative is true if the test should be "not zero".
+  /// This will also optimize the test bit instruction when possible.
+  MachineInstr *emitTestBit(Register TestReg, uint64_t Bit, bool IsNegative,
+                            MachineBasicBlock *DstMBB,
+                            MachineIRBuilder &MIB) const;
+
   // Equivalent to the i32shift_a and friends from AArch64InstrInfo.td.
   // We use these manually instead of using the importer since it doesn't
   // support SDNodeXForm.
   ComplexRendererFns selectShiftA_32(const MachineOperand &Root) const;
   ComplexRendererFns selectShiftB_32(const MachineOperand &Root) const;
   ComplexRendererFns selectShiftA_64(const MachineOperand &Root) const;
   ComplexRendererFns selectShiftB_64(const MachineOperand &Root) const;
 
while (MachineInstr *MI = getDefIgnoringCopies(Reg, MRI)) {
     if (!NextReg.isValid())
       return Reg;
     Reg = NextReg;
   }
 
   return Reg;
 }
 
+MachineInstr *AArch64InstructionSelector::emitTestBit(
+    Register TestReg, uint64_t Bit, bool IsNegative, MachineBasicBlock *DstMBB,
+    MachineIRBuilder &MIB) const {
+  MachineRegisterInfo &MRI = *MIB.getMRI();
+#ifndef NDEBUG
+  assert(ProduceNonFlagSettingCondBr &&
+         "Cannot emit TB(N)Z with speculation tracking!");
+  assert(TestReg.isValid());
+  LLT Ty = MRI.getType(TestReg);
+  unsigned Size = Ty.getSizeInBits();
+  assert(Bit < Size &&
+         "Bit to test must be smaler than the size of a test register!");
+  assert(Ty.isScalar() && "Expected a scalar!");
+  assert(Size >= 32 && "Expected at least a 32-bit register!");
+#endif
+
+  // Attempt to optimize the test bit by walking over instructions.
+  TestReg = getTestBitReg(TestReg, Bit, IsNegative, MRI);
+  bool UseWReg = Bit < 32;
+
+  // When the test register is a 64-bit register, we have to narrow to make
+  // TBNZW work.
+  if (UseWReg)
+    TestReg = narrowExtendRegIfNeeded(TestReg, MIB);
+
+  static const unsigned OpcTable[2][2] = {{AArch64::TBZX, AArch64::TBNZX},
+                                          {AArch64::TBZW, AArch64::TBNZW}};
+  unsigned Opc = OpcTable[UseWReg][IsNegative];
+  auto TestBitMI =
+      MIB.buildInstr(Opc).addReg(TestReg).addImm(Bit).addMBB(DstMBB);
+  constrainSelectedInstRegOperands(*TestBitMI, TII, TRI, RBI);
+  return &*TestBitMI;
+}
+
 bool AArch64InstructionSelector::tryOptAndIntoCompareBranch(
     MachineInstr *AndInst, int64_t CmpConstant, const CmpInst::Predicate &Pred,
     MachineBasicBlock *DstMBB, MachineIRBuilder &MIB) const {
   // Given something like this:
   //
   //  %x = ...Something...
   //  %one = G_CONSTANT i64 1
   //  %zero = G_CONSTANT i64 0
bool AArch64InstructionSelector::tryOptAndIntoCompareBranch(
   // Check if the AND has a constant on its RHS which we can use as a mask.
   // If it's a power of 2, then it's the same as checking a specific bit.
   // (e.g, ANDing with 8 == ANDing with 000...100 == testing if bit 3 is set)
   auto MaybeBit =
       getConstantVRegValWithLookThrough(AndInst->getOperand(2).getReg(), MRI);
   if (!MaybeBit || !isPowerOf2_64(MaybeBit->Value))
     return false;
 
-  // Try to optimize the TB(N)Z.
   uint64_t Bit = Log2_64(static_cast<uint64_t>(MaybeBit->Value));
   Register TestReg = AndInst->getOperand(1).getReg();
   bool Invert = Pred == CmpInst::Predicate::ICMP_NE;
-  TestReg = getTestBitReg(TestReg, Bit, Invert, MRI);
 
-  // Choose the correct TB(N)Z opcode to use.
-  unsigned Opc = 0;
-  if (Bit < 32) {
-    // When the bit is less than 32, we have to use a TBZW even if we're on a 64
-    // bit register.
-    Opc = Invert ? AArch64::TBNZW : AArch64::TBZW;
-    TestReg = narrowExtendRegIfNeeded(TestReg, MIB);
-  } else {
-    // Same idea for when Bit >= 32. We don't have to narrow here, because if
-    // Bit > 32, then the G_CONSTANT must be outside the range of valid 32-bit
-    // values. So, we must have a s64.
-    Opc = Invert ? AArch64::TBNZX : AArch64::TBZX;
-  }
-
-  // Construct the branch.
-  auto BranchMI =
-      MIB.buildInstr(Opc).addReg(TestReg).addImm(Bit).addMBB(DstMBB);
-  constrainSelectedInstRegOperands(*BranchMI, TII, TRI, RBI);
+  // Emit a TB(N)Z.
+  emitTestBit(TestReg, Bit, Invert, DstMBB, MIB);
   return true;
 }
 
 bool AArch64InstructionSelector::selectCompareBranch(
     MachineInstr &I, MachineFunction &MF, MachineRegisterInfo &MRI) const {
 
   const Register CondReg = I.getOperand(0).getReg();
   MachineBasicBlock *DestMBB = I.getOperand(1).getMBB();