[GlobalISel][Utils] Optionally accept alternate COPY opcodes in getDefSrcRegIgnoringCopies

[jackoalan]

Summary

Adds optional argument to provide a set of alternate COPY opcodes for
getDefSrcRegIgnoringCopies to recognize. This is useful for instruction
selection to look through predecessors that will be directly selected as COPY
once iterated to.

[jackoalan]

Fix clang-format issue

[jackoalan]

Abandoning this for now. A more sensible approach to the use case has been discovered.

llvm/include/llvm/CodeGen/GlobalISel/Utils.h

struct DefinitionAndSourceRegister {
   MachineInstr *MI;
   Register Reg;
 };
 
 /// Find the def instruction for \p Reg, and underlying value Register folding
 /// away any copies.
 ///
 /// Also walks through hints such as G_ASSERT_ZEXT.
-Optional<DefinitionAndSourceRegister>
-getDefSrcRegIgnoringCopies(Register Reg, const MachineRegisterInfo &MRI);
+///
+/// Optionally, an alternative set of copy opcodes can be recognized through
+/// \p CopyOpcodes.
+Optional<DefinitionAndSourceRegister> getDefSrcRegIgnoringCopies(
+    Register Reg, const MachineRegisterInfo &MRI,
+    ArrayRef<unsigned> CopyOpcodes = {TargetOpcode::COPY});
 
 /// Find the def instruction for \p Reg, folding away any trivial copies. May
 /// return nullptr if \p Reg is not a generic virtual register.
 ///
 /// Also walks through hints such as G_ASSERT_ZEXT.
-MachineInstr *getDefIgnoringCopies(Register Reg,
-                                   const MachineRegisterInfo &MRI);
+///
+/// Optionally, an alternative set of copy opcodes can be recognized through
+/// \p CopyOpcodes.
+MachineInstr *getDefIgnoringCopies(Register Reg, const MachineRegisterInfo &MRI,
+                                   ArrayRef<unsigned> CopyOpcodes = {
+                                       TargetOpcode::COPY});
 
 /// Find the source register for \p Reg, folding away any trivial copies. It
 /// will be an output register of the instruction that getDefIgnoringCopies
 /// returns. May return an invalid register if \p Reg is not a generic virtual
 /// register.
 ///
 /// Also walks through hints such as G_ASSERT_ZEXT.
-Register getSrcRegIgnoringCopies(Register Reg, const MachineRegisterInfo &MRI);
+///
+/// Optionally, an alternative set of copy opcodes can be recognized through
+/// \p CopyOpcodes.
+Register getSrcRegIgnoringCopies(Register Reg, const MachineRegisterInfo &MRI,
+                                 ArrayRef<unsigned> CopyOpcodes = {
+                                     TargetOpcode::COPY});
 
 // Templated variant of getOpcodeDef returning a MachineInstr derived T.
 /// See if Reg is defined by an single def instruction of type T
 /// Also try to do trivial folding if it's a COPY with
 /// same types. Returns null otherwise.
 template <class T>
 T *getOpcodeDef(Register Reg, const MachineRegisterInfo &MRI) {
   MachineInstr *DefMI = getDefIgnoringCopies(Reg, MRI);

llvm/lib/CodeGen/GlobalISel/Utils.cpp

 llvm::getConstantFPVRegVal(Register VReg, const MachineRegisterInfo &MRI) {
   MachineInstr *MI = MRI.getVRegDef(VReg);
   if (TargetOpcode::G_FCONSTANT != MI->getOpcode())
     return nullptr;
   return MI->getOperand(1).getFPImm();
 }
 
 Optional<DefinitionAndSourceRegister>
-llvm::getDefSrcRegIgnoringCopies(Register Reg, const MachineRegisterInfo &MRI) {
+llvm::getDefSrcRegIgnoringCopies(Register Reg, const MachineRegisterInfo &MRI,
+                                 ArrayRef<unsigned> CopyOpcodes) {
   Register DefSrcReg = Reg;
   auto *DefMI = MRI.getVRegDef(Reg);
   auto DstTy = MRI.getType(DefMI->getOperand(0).getReg());
   if (!DstTy.isValid())
     return None;
   unsigned Opc = DefMI->getOpcode();
-  while (Opc == TargetOpcode::COPY || isPreISelGenericOptimizationHint(Opc)) {
+  while (is_contained(CopyOpcodes, Opc) ||
+         isPreISelGenericOptimizationHint(Opc)) {
     Register SrcReg = DefMI->getOperand(1).getReg();
     auto SrcTy = MRI.getType(SrcReg);
     if (!SrcTy.isValid())
       break;
     DefMI = MRI.getVRegDef(SrcReg);
     DefSrcReg = SrcReg;
     Opc = DefMI->getOpcode();
   }
   return DefinitionAndSourceRegister{DefMI, DefSrcReg};
 }
 
 MachineInstr *llvm::getDefIgnoringCopies(Register Reg,
-                                         const MachineRegisterInfo &MRI) {
+                                         const MachineRegisterInfo &MRI,
+                                         ArrayRef<unsigned> CopyOpcodes) {
   Optional<DefinitionAndSourceRegister> DefSrcReg =
-      getDefSrcRegIgnoringCopies(Reg, MRI);
+      getDefSrcRegIgnoringCopies(Reg, MRI, CopyOpcodes);
   return DefSrcReg ? DefSrcReg->MI : nullptr;
 }
 
 Register llvm::getSrcRegIgnoringCopies(Register Reg,
-                                       const MachineRegisterInfo &MRI) {
+                                       const MachineRegisterInfo &MRI,
+                                       ArrayRef<unsigned> CopyOpcodes) {
   Optional<DefinitionAndSourceRegister> DefSrcReg =
-      getDefSrcRegIgnoringCopies(Reg, MRI);
+      getDefSrcRegIgnoringCopies(Reg, MRI, CopyOpcodes);
   return DefSrcReg ? DefSrcReg->Reg : Register();
 }
 
 MachineInstr *llvm::getOpcodeDef(unsigned Opcode, Register Reg,
                                  const MachineRegisterInfo &MRI) {
   MachineInstr *DefMI = getDefIgnoringCopies(Reg, MRI);
   return DefMI && DefMI->getOpcode() == Opcode ? DefMI : nullptr;
 }

llvm/unittests/CodeGen/GlobalISel/GISelUtilsTest.cpp

 //===- GISelUtilsTest.cpp -------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "GISelMITest.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
 static const LLT S1 = LLT::scalar(1);
 static const LLT S8 = LLT::scalar(8);
TEST(GISelUtilsTest, getLCMType) {
 
   EXPECT_EQ(V2S64, getLCMType(V2S64, P0));
   EXPECT_EQ(V2P0, getLCMType(P0, V2S64));
 
   EXPECT_EQ(V2S64, getLCMType(V2S64, P1));
   EXPECT_EQ(V4P1, getLCMType(P1, V2S64));
 }
 
+TEST_F(AArch64GISelMITest, AltIgnoreCopies) {
+  setUp();
+  if (!TM)
+    return;
+
+  auto X = B.buildInstr(TargetOpcode::G_IMPLICIT_DEF, {S1}, {});
+  auto Freeze = B.buildInstr(TargetOpcode::G_FREEZE, {S1}, {X});
+  auto Copy = B.buildInstr(TargetOpcode::COPY, {S1}, {Freeze});
+
+  auto DefSrcReg =
+      getDefSrcRegIgnoringCopies(Copy->getOperand(0).getReg(), *MRI,
+                                 {TargetOpcode::COPY, TargetOpcode::G_FREEZE});
+
+  EXPECT_EQ(DefSrcReg->MI, X);
+  EXPECT_EQ(DefSrcReg->Reg, X->getOperand(0).getReg());
 }
+
+} // namespace