[AArch64][GlobalISel] Support shifted register form in emitTST

[paquette]

Summary

Support emitting ANDSXrs and ANDSWrs in emitTST. Update opt-fold-compare.mir to show that it works.

[aemerson]

Add a test for ANDSWrs?

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

MachineInstr *emitInstr(unsigned Opcode,
                           const ComplexRendererFns &RenderFns = None) const;
   MachineInstr *emitADD(Register DefReg, MachineOperand &LHS,
                         MachineOperand &RHS,
                         MachineIRBuilder &MIRBuilder) const;
   MachineInstr *emitADDS(Register Dst, MachineOperand &LHS, MachineOperand &RHS,
                          MachineIRBuilder &MIRBuilder) const;
   MachineInstr *emitCMN(MachineOperand &LHS, MachineOperand &RHS,
                         MachineIRBuilder &MIRBuilder) const;
-  MachineInstr *emitTST(const Register &LHS, const Register &RHS,
+  MachineInstr *emitTST(MachineOperand &LHS, MachineOperand &RHS,
                         MachineIRBuilder &MIRBuilder) const;
   MachineInstr *emitExtractVectorElt(Optional<Register> DstReg,
                                      const RegisterBank &DstRB, LLT ScalarTy,
                                      Register VecReg, unsigned LaneIdx,
                                      MachineIRBuilder &MIRBuilder) const;
 
   /// Helper function for selecting G_FCONSTANT. If the G_FCONSTANT can be
   /// materialized using a FMOV instruction, then update MI and return it.

 AArch64InstructionSelector::emitCMN(MachineOperand &LHS, MachineOperand &RHS,
                                     MachineIRBuilder &MIRBuilder) const {
   MachineRegisterInfo &MRI = MIRBuilder.getMF().getRegInfo();
   bool Is32Bit = (MRI.getType(LHS.getReg()).getSizeInBits() == 32);
   return emitADDS(Is32Bit ? AArch64::WZR : AArch64::XZR, LHS, RHS, MIRBuilder);
 }
 
 MachineInstr *
-AArch64InstructionSelector::emitTST(const Register &LHS, const Register &RHS,
+AArch64InstructionSelector::emitTST(MachineOperand &LHS, MachineOperand &RHS,
                                     MachineIRBuilder &MIRBuilder) const {
+  assert(LHS.isReg() && RHS.isReg() && "Expected register operands?");
   MachineRegisterInfo &MRI = MIRBuilder.getMF().getRegInfo();
-  unsigned RegSize = MRI.getType(LHS).getSizeInBits();
+  unsigned RegSize = MRI.getType(LHS.getReg()).getSizeInBits();
   bool Is32Bit = (RegSize == 32);
-  static const unsigned OpcTable[2][2]{{AArch64::ANDSXrr, AArch64::ANDSXri},
-                                       {AArch64::ANDSWrr, AArch64::ANDSWri}};
+  const unsigned OpcTable[3][2] = {{AArch64::ANDSXri, AArch64::ANDSWri},
+                                   {AArch64::ANDSXrs, AArch64::ANDSWrs},
+                                   {AArch64::ANDSXrr, AArch64::ANDSWrr}};
   Register ZReg = Is32Bit ? AArch64::WZR : AArch64::XZR;
-
-  // We might be able to fold in an immediate into the TST. We need to make sure
-  // it's a logical immediate though, since ANDS requires that.
-  auto ValAndVReg = getConstantVRegValWithLookThrough(RHS, MRI);
-  bool IsImmForm = ValAndVReg.hasValue() &&
-                   AArch64_AM::isLogicalImmediate(ValAndVReg->Value, RegSize);
-  unsigned Opc = OpcTable[Is32Bit][IsImmForm];
-  auto TstMI = MIRBuilder.buildInstr(Opc, {ZReg}, {LHS});
-
-  if (IsImmForm)
+  // ANDS needs a logical immediate for its immediate form. Check if we can
+  // fold one in.
+  if (auto ValAndVReg = getConstantVRegValWithLookThrough(RHS.getReg(), MRI)) {
+    if (AArch64_AM::isLogicalImmediate(ValAndVReg->Value, RegSize)) {
+      auto TstMI = MIRBuilder.buildInstr(OpcTable[0][Is32Bit], {ZReg}, {LHS});
       TstMI.addImm(
           AArch64_AM::encodeLogicalImmediate(ValAndVReg->Value, RegSize));
-  else
-    TstMI.addUse(RHS);
-
       constrainSelectedInstRegOperands(*TstMI, TII, TRI, RBI);
       return &*TstMI;
     }
+  }
+
+  if (auto Fns = selectLogicalShiftedRegister(RHS))
+    return emitInstr(OpcTable[1][Is32Bit], {ZReg}, {LHS}, MIRBuilder, Fns);
+  return emitInstr(OpcTable[2][Is32Bit], {ZReg}, {LHS, RHS}, MIRBuilder);
+}
 
 std::pair<MachineInstr *, CmpInst::Predicate>
 AArch64InstructionSelector::emitIntegerCompare(
     MachineOperand &LHS, MachineOperand &RHS, MachineOperand &Predicate,
     MachineIRBuilder &MIRBuilder) const {
   assert(LHS.isReg() && RHS.isReg() && "Expected LHS and RHS to be registers!");
   assert(Predicate.isPredicate() && "Expected predicate?");
   MachineRegisterInfo &MRI = MIRBuilder.getMF().getRegInfo();
MachineInstr *AArch64InstructionSelector::tryFoldIntegerCompare(
   // tst x, y
   if (!isUnsignedICMPPred(P) && LHSDef &&
       LHSDef->getOpcode() == TargetOpcode::G_AND) {
     // Make sure that the RHS is 0.
     auto ValAndVReg = getConstantVRegValWithLookThrough(RHS.getReg(), MRI);
     if (!ValAndVReg || ValAndVReg->Value != 0)
       return nullptr;
 
-    return emitTST(LHSDef->getOperand(1).getReg(),
-                   LHSDef->getOperand(2).getReg(), MIRBuilder);
+    return emitTST(LHSDef->getOperand(1),
+                   LHSDef->getOperand(2), MIRBuilder);
   }
 
   return nullptr;
 }
 
 MachineInstr *AArch64InstructionSelector::tryOptArithImmedIntegerCompare(
     MachineOperand &LHS, MachineOperand &RHS, CmpInst::Predicate &P,
     MachineIRBuilder &MIB) const {

llvm/test/CodeGen/AArch64/GlobalISel/opt-fold-compare.mir

bb.0:
     %0:gpr(s64) = COPY $x0
     %1:gpr(s64) = COPY $x1
     ; When we find the defs of the LHS and RHS of the compare, we walk over
     ; copies. Make sure that we don't crash when we hit a copy from a physical
     ; register.
     %7:gpr(s32) = G_ICMP intpred(eq), %0, %1
     $w0 = COPY %7(s32)
     RET_ReallyLR implicit $x0
+
+...
+---
+name:            tst_fold_shift
+alignment:       4
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; We should fold the G_SHL into the ANDS to get ANDSXrs.
+    ;
+    ; CHECK-LABEL: name: tst_fold_shift
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: %copy:gpr64 = COPY $x1
+    ; CHECK: %zero:gpr64 = COPY $xzr
+    ; CHECK: [[MOVi32imm:%[0-9]+]]:gpr32 = MOVi32imm 1
+    ; CHECK: %one:gpr64 = SUBREG_TO_REG 0, [[MOVi32imm]], %subreg.sub_32
+    ; CHECK: $xzr = ANDSXrs %zero, %copy, 16, implicit-def $nzcv
+    ; CHECK: %select:gpr64 = CSELXr %one, %zero, 0, implicit $nzcv
+    ; CHECK: $x0 = COPY %select
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:gpr(s64) = COPY $x1
+    %zero:gpr(s64) = G_CONSTANT i64 0
+    %one:gpr(s64) = G_CONSTANT i64 1
+    %cst:gpr(s64) = G_CONSTANT i64 16
+    %shift:gpr(s64) = G_SHL %copy(s64), %cst(s64)
+    %and:gpr(s64) = G_AND %zero, %shift
+    %cmp:gpr(s32) = G_ICMP intpred(eq), %and(s64), %zero
+    %cmp_trunc:gpr(s1) = G_TRUNC %cmp(s32)
+    %select:gpr(s64) = G_SELECT %cmp_trunc(s1), %one, %zero
+    $x0 = COPY %select(s64)
+    RET_ReallyLR implicit $x0