[GlobalISel] Allow the ArtifactValueFinder to return the best availab…

…le register on failure.

In some cases, like with inserts, we may have a matching size register already,
but still decide to try to look further. This change adds a CurrentBest
register to the value finder state, and any time a method fails to make progress,
returns that register (which may just be an empty Register).

To facilitate this, add a new entry point to the findValueFromDef() function
which initializes this state.

Also fix the build vector finder to return the current build_vector if all
sources are being requested.

Differential Revision: https://reviews.llvm.org/D107017

llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h

@@ -23,6 +23,7 @@
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/Register.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "legalizer"
@@ -545,12 +546,14 @@ class LegalizationArtifactCombiner {
     MachineIRBuilder &MIB;
     const LegalizerInfo &LI;
 
-  private:
+    // Stores the best register found in the current query so far.
+    Register CurrentBest = Register();
+
     /// Given an concat_vector op \p Concat and a start bit and size, try to
     /// find the origin of the value defined by that start position and size.
     ///
-    /// \returns A register if a value can be found, otherwise an empty
-    /// Register.
+    /// \returns a register with the requested size, or the current best
+    /// register found during the current query.
     Register findValueFromConcat(GConcatVectors &Concat, unsigned StartBit,
                                  unsigned Size) {
       assert(Size > 0);
@@ -567,22 +570,22 @@ class LegalizationArtifactCombiner {
       // FIXME: we might be able return multiple sources? Or create an
       // appropriate concat to make it fit.
       if (InRegOffset + Size > SrcSize)
-        return Register();
+        return CurrentBest;
 
-      // If the bits exactly cover a single source, then return the operand as
-      // our value reg.
       Register SrcReg = Concat.getReg(StartSrcIdx);
-      if (InRegOffset == 0 && Size == SrcSize)
-        return SrcReg; // A source operand matches exactly.
+      if (InRegOffset == 0 && Size == SrcSize) {
+        CurrentBest = SrcReg;
+        return findValueFromDefImpl(SrcReg, 0, Size);
+      }
 
-      return findValueFromDef(SrcReg, InRegOffset, Size);
+      return findValueFromDefImpl(SrcReg, InRegOffset, Size);
     }
 
     /// Given an build_vector op \p BV and a start bit and size, try to find
     /// the origin of the value defined by that start position and size.
     ///
-    /// \returns A register if a value can be found, otherwise an empty
-    /// Register.
+    /// \returns a register with the requested size, or the current best
+    /// register found during the current query.
     Register findValueFromBuildVector(GBuildVector &BV, unsigned StartBit,
                                       unsigned Size) {
       assert(Size > 0);
@@ -597,25 +600,29 @@ class LegalizationArtifactCombiner {
       unsigned InRegOffset = StartBit % SrcSize;
 
       if (InRegOffset != 0)
-        return Register(); // Give up, bits don't start at a scalar source.
+        return CurrentBest; // Give up, bits don't start at a scalar source.
       if (Size < SrcSize)
-        return Register(); // Scalar source is too large for requested bits.
+        return CurrentBest; // Scalar source is too large for requested bits.
 
       // If the bits cover multiple sources evenly, then create a new
       // build_vector to synthesize the required size, if that's been requested.
       if (Size > SrcSize) {
         if (Size % SrcSize > 0)
-          return Register(); // Isn't covered exactly by sources.
+          return CurrentBest; // Isn't covered exactly by sources.
 
         unsigned NumSrcsUsed = Size / SrcSize;
+        // If we're requesting all of the sources, just return this def.
+        if (NumSrcsUsed == BV.getNumSources())
+          return BV.getReg(0);
+
         LLT SrcTy = MRI.getType(Src1Reg);
         LLT NewBVTy = LLT::fixed_vector(NumSrcsUsed, SrcTy);
 
         // Check if the resulting build vector would be legal.
         LegalizeActionStep ActionStep =
             LI.getAction({TargetOpcode::G_BUILD_VECTOR, {NewBVTy, SrcTy}});
         if (ActionStep.Action != LegalizeActions::Legal)
-          return Register();
+          return CurrentBest;
 
         SmallVector<Register> NewSrcs;
         for (unsigned SrcIdx = StartSrcIdx; SrcIdx < StartSrcIdx + NumSrcsUsed;
@@ -631,8 +638,8 @@ class LegalizationArtifactCombiner {
     /// Given an G_INSERT op \p MI and a start bit and size, try to find
     /// the origin of the value defined by that start position and size.
     ///
-    /// \returns A register if a value can be found, otherwise an empty
-    /// Register.
+    /// \returns a register with the requested size, or the current best
+    /// register found during the current query.
     Register findValueFromInsert(MachineInstr &MI, unsigned StartBit,
                                  unsigned Size) {
       assert(MI.getOpcode() == TargetOpcode::G_INSERT);
@@ -686,28 +693,25 @@ class LegalizationArtifactCombiner {
       if (EndBit <= InsertOffset || InsertedEndBit <= StartBit) {
         SrcRegToUse = ContainerSrcReg;
         NewStartBit = StartBit;
-        return findValueFromDef(SrcRegToUse, NewStartBit, Size);
+        return findValueFromDefImpl(SrcRegToUse, NewStartBit, Size);
       }
       if (InsertOffset <= StartBit && EndBit <= InsertedEndBit) {
         SrcRegToUse = InsertedReg;
         NewStartBit = StartBit - InsertOffset;
-        return findValueFromDef(SrcRegToUse, NewStartBit, Size);
+        if (NewStartBit == 0 &&
+            Size == MRI.getType(SrcRegToUse).getSizeInBits())
+          CurrentBest = SrcRegToUse;
+        return findValueFromDefImpl(SrcRegToUse, NewStartBit, Size);
       }
       // The bit range spans both the inserted and container regions.
       return Register();
     }
 
-  public:
-    ArtifactValueFinder(MachineRegisterInfo &Mri, MachineIRBuilder &Builder,
-                        const LegalizerInfo &Info)
-        : MRI(Mri), MIB(Builder), LI(Info) {}
-
-    /// Try to find a source of the value defined in the def \p DefReg, starting
-    /// at position \p StartBit with size \p Size.
-    /// \returns an empty Register if no value could be found, or \p DefReg if
-    /// if that was the best we could do.
-    Register findValueFromDef(Register DefReg, unsigned StartBit,
-                              unsigned Size) {
+    /// Internal implementation for findValueFromDef(). findValueFromDef()
+    /// initializes some data like the CurrentBest register, which this method
+    /// and its callees rely upon.
+    Register findValueFromDefImpl(Register DefReg, unsigned StartBit,
+                                  unsigned Size) {
       MachineInstr *Def = getDefIgnoringCopies(DefReg, MRI);
       // If the instruction has a single def, then simply delegate the search.
       // For unmerge however with multiple defs, we need to compute the offset
@@ -725,26 +729,42 @@ class LegalizationArtifactCombiner {
         }
         Register SrcReg = Def->getOperand(Def->getNumOperands() - 1).getReg();
         Register SrcOriginReg =
-            findValueFromDef(SrcReg, StartBit + DefStartBit, Size);
+            findValueFromDefImpl(SrcReg, StartBit + DefStartBit, Size);
         if (SrcOriginReg)
           return SrcOriginReg;
         // Failed to find a further value. If the StartBit and Size perfectly
         // covered the requested DefReg, return that since it's better than
         // nothing.
         if (StartBit == 0 && Size == DefSize)
           return DefReg;
-        return Register();
+        return CurrentBest;
       }
       case TargetOpcode::G_BUILD_VECTOR:
         return findValueFromBuildVector(cast<GBuildVector>(*Def), StartBit,
                                         Size);
       case TargetOpcode::G_INSERT:
         return findValueFromInsert(*Def, StartBit, Size);
       default:
-        return Register();
+        return CurrentBest;
       }
     }
 
+  public:
+    ArtifactValueFinder(MachineRegisterInfo &Mri, MachineIRBuilder &Builder,
+                        const LegalizerInfo &Info)
+        : MRI(Mri), MIB(Builder), LI(Info) {}
+
+    /// Try to find a source of the value defined in the def \p DefReg, starting
+    /// at position \p StartBit with size \p Size.
+    /// \returns a register with the requested size, or an empty Register if no
+    /// better value could be found.
+    Register findValueFromDef(Register DefReg, unsigned StartBit,
+                              unsigned Size) {
+      CurrentBest = Register();
+      Register FoundReg = findValueFromDefImpl(DefReg, StartBit, Size);
+      return FoundReg != DefReg ? FoundReg : Register();
+    }
+
     /// Try to combine the defs of an unmerge \p MI by attempting to find
     /// values that provides the bits for each def reg.
     /// \returns true if all the defs of the unmerge have been made dead.
@@ -760,9 +780,8 @@ class LegalizationArtifactCombiner {
           DeadDefs[DefIdx] = true;
           continue;
         }
-        Register FoundVal =
-            findValueFromDef(DefReg, 0, DestTy.getSizeInBits());
-        if (!FoundVal || FoundVal == DefReg)
+        Register FoundVal = findValueFromDef(DefReg, 0, DestTy.getSizeInBits());
+        if (!FoundVal)
           continue;
         if (MRI.getType(FoundVal) != DestTy)
           continue;