[GlobalISel][AArch64][NFC] Use getDefIgnoringCopies from Utils where we can

[paquette]

Summary

There are a few places where we walk over copies throughout AArch64InstructionSelector.cpp. In Utils, there's a function that does exactly this which we can use instead.

Note that the utility function works with the case where we run into a COPY from a physical register. We've run into bugs with this a couple times, so using it should defend us from similar future bugs.

[arsenm]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

3239–3240

I'm not sure getDefIgnoringCopies handles physical registers correctly, does that have a test somewhere?

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

3239–3240

It should. It gets the source type using MRI.getType, which should return LLT {} when you give it a physical register. This should make the source type invalid, which ends the search.

I'll add a test anyway though, since I'm not sure if one exists.

[paquette]

Add a testcase showing that we won't crash if we encounter a copy from a physical register.

[aemerson]

LGTM.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

void AArch64InstructionSelector::collectShuffleMaskIndices(
     MachineInstr &I, MachineRegisterInfo &MRI,
     SmallVectorImpl<Optional<int>> &Idxs) const {
   MachineInstr *MaskDef = MRI.getVRegDef(I.getOperand(3).getReg());
   assert(
       MaskDef->getOpcode() == TargetOpcode::G_BUILD_VECTOR &&
       "G_SHUFFLE_VECTOR should have a constant mask operand as G_BUILD_VECTOR");
   // Find the constant indices.
   for (unsigned i = 1, e = MaskDef->getNumOperands(); i < e; ++i) {
-    MachineInstr *ScalarDef = MRI.getVRegDef(MaskDef->getOperand(i).getReg());
-    assert(ScalarDef && "Could not find vreg def of shufflevec index op");
     // Look through copies.
-    while (ScalarDef->getOpcode() == TargetOpcode::COPY) {
-      ScalarDef = MRI.getVRegDef(ScalarDef->getOperand(1).getReg());
-      assert(ScalarDef && "Could not find def of copy operand");
-    }
+    MachineInstr *ScalarDef =
+        getDefIgnoringCopies(MaskDef->getOperand(i).getReg(), MRI);
+    assert(ScalarDef && "Could not find vreg def of shufflevec index op");
     if (ScalarDef->getOpcode() != TargetOpcode::G_CONSTANT) {
       // This be an undef if not a constant.
       assert(ScalarDef->getOpcode() == TargetOpcode::G_IMPLICIT_DEF);
       Idxs.push_back(None);
     } else {
       Idxs.push_back(ScalarDef->getOperand(1).getCImm()->getSExtValue());
     }
   }
MachineInstr *AArch64InstructionSelector::tryFoldIntegerCompare(
   // x = G_SUB 0, y
   // G_ICMP z, x
   //
   // In this case, we can fold the G_SUB into the G_ICMP using a CMN instead.
   // e.g:
   //
   // cmn z, y
 
-  // Helper lambda to find the def.
-  auto FindDef = [&](Register VReg) {
-    MachineInstr *Def = MRI.getVRegDef(VReg);
-    while (Def) {
-      if (Def->getOpcode() != TargetOpcode::COPY)
-        break;
-      // Copies can be from physical registers. If we hit this, we're done.
-      if (TargetRegisterInfo::isPhysicalRegister(Def->getOperand(1).getReg()))
-        break;
-      Def = MRI.getVRegDef(Def->getOperand(1).getReg());
-    }
-    return Def;
-  };
-
   // Helper lambda to detect the subtract followed by the compare.
   // Takes in the def of the LHS or RHS, and checks if it's a subtract from 0.
   auto IsCMN = [&](MachineInstr *DefMI, const AArch64CC::CondCode &CC) {
     if (!DefMI || DefMI->getOpcode() != TargetOpcode::G_SUB)
       return false;
 
     // Need to make sure NZCV is the same at the end of the transformation.
     if (CC != AArch64CC::EQ && CC != AArch64CC::NE)
auto IsCMN = [&](MachineInstr *DefMI, const AArch64CC::CondCode &CC) {
     if (!ValAndVReg || ValAndVReg->Value != 0)
       return false;
 
     // This can safely be represented as a CMN.
     return true;
   };
 
   // Check if the RHS or LHS of the G_ICMP is defined by a SUB
-  MachineInstr *LHSDef = FindDef(LHS.getReg());
-  MachineInstr *RHSDef = FindDef(RHS.getReg());
+  MachineInstr *LHSDef = getDefIgnoringCopies(LHS.getReg(), MRI);
+  MachineInstr *RHSDef = getDefIgnoringCopies(RHS.getReg(), MRI);
   CmpInst::Predicate P = (CmpInst::Predicate)Predicate.getPredicate();
   const AArch64CC::CondCode CC = changeICMPPredToAArch64CC(P);
 
   // Given this:
   //
   // x = G_SUB 0, y
   // G_ICMP x, z
   //

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

bb.0:
     ; This immediate can't be represented as a logical immediate. We shouldn't
     ; select ANDSWri.
     %4:gpr(s32) = G_CONSTANT i32 -1
 
     %5:gpr(s32) = G_AND %1, %4
     %6:gpr(s32) = G_ICMP intpred(eq), %5(s32), %2
     $w0 = COPY %6(s32)
     RET_ReallyLR implicit $w0
+
+...
+---
+name:            test_physreg_copy
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: test_physreg_copy
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK: $xzr = SUBSXrr [[COPY]], [[COPY1]], implicit-def $nzcv
+    ; CHECK: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr $wzr, $wzr, 1, implicit $nzcv
+    ; CHECK: $w0 = COPY [[CSINCWr]]
+    ; CHECK: RET_ReallyLR implicit $x0
+    %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