[PeepholeOptimizer] Refactor optimizeUncoalescable logic. NFC

[bruno]

Summary

The PeepholeOptimizer misses some rewrite opportunities since it can't look through PHI instructions.
This patch cleans the current logic a bit while preparing for another patch that can handle PHIs. Changes include:

• Create a new CopyRewriter for Uncoalescable copy-like instructions.
• Change the ValueTracker to return a ValueTrackerResult.

This makes optimizeUncoalescable looks more like optimizeCoalescable and use the CopyRewritter infrastructure.

[qcolombet]

Hi Bruno,

Nice refactoring!

LGTM with one nit.
Please commit with that updated.

Cheers,
-Quentin

lib/CodeGen/PeepholeOptimizer.cpp

339 ↗

(On Diff #29705)

The comment here needs further update as SrcReg and SrcSubReg are not parameters anymore.

[bruno]

Good catch, thanks Quentin!

llvm/trunk/lib/CodeGen/PeepholeOptimizer.cpp

 STATISTIC(NumCmps,       "Number of compares eliminated");
 STATISTIC(NumImmFold,    "Number of move immediate folded");
 STATISTIC(NumLoadFold,   "Number of loads folded");
 STATISTIC(NumSelects,    "Number of selects optimized");
 STATISTIC(NumUncoalescableCopies, "Number of uncoalescable copies optimized");
 STATISTIC(NumRewrittenCopies, "Number of copies rewritten");
 
 namespace {
+  class ValueTrackerResult;
+
   class PeepholeOptimizer : public MachineFunctionPass {
     const TargetInstrInfo *TII;
     const TargetRegisterInfo *TRI;
     MachineRegisterInfo   *MRI;
     MachineDominatorTree  *DT;  // Machine dominator tree
 
   public:
     static char ID; // Pass identification
private:
     bool isUncoalescableCopy(const MachineInstr &MI) {
       return MI.isBitcast() ||
              (!DisableAdvCopyOpt &&
               (MI.isRegSequenceLike() || MI.isInsertSubregLike() ||
                MI.isExtractSubregLike()));
     }
   };
 
+  /// \brief Helper class to hold a reply for ValueTracker queries. Contains the
+  /// returned sources for a given search and the instructions where the sources
+  /// were tracked from.
+  class ValueTrackerResult {
+  private:
+    /// Track all sources found by one ValueTracker query.
+    SmallVector<TargetInstrInfo::RegSubRegPair, 2> RegSrcs;
+
+    /// Instruction using the sources in 'RegSrcs'.
+    const MachineInstr *Inst;
+
+  public:
+    ValueTrackerResult() : Inst(nullptr) {}
+    ValueTrackerResult(unsigned Reg, unsigned SubReg) : Inst(nullptr) {
+      addSource(Reg, SubReg);
+    }
+
+    bool isValid() const { return getNumSources() > 0; }
+
+    void setInst(const MachineInstr *I) { Inst = I; }
+    const MachineInstr *getInst() const { return Inst; }
+
+    void clear() {
+      RegSrcs.clear();
+      Inst = nullptr;
+    }
+
+    void addSource(unsigned SrcReg, unsigned SrcSubReg) {
+      RegSrcs.push_back(TargetInstrInfo::RegSubRegPair(SrcReg, SrcSubReg));
+    }
+
+    void setSource(int Idx, unsigned SrcReg, unsigned SrcSubReg) {
+      assert(Idx < getNumSources() && "Reg pair source out of index");
+      RegSrcs[Idx] = TargetInstrInfo::RegSubRegPair(SrcReg, SrcSubReg);
+    }
+
+    int getNumSources() const { return RegSrcs.size(); }
+
+    unsigned getSrcReg(int Idx) const {
+      assert(Idx < getNumSources() && "Reg source out of index");
+      return RegSrcs[Idx].Reg;
+    }
+
+    unsigned getSrcSubReg(int Idx) const {
+      assert(Idx < getNumSources() && "SubReg source out of index");
+      return RegSrcs[Idx].SubReg;
+    }
+  };
+
   /// \brief Helper class to track the possible sources of a value defined by
   /// a (chain of) copy related instructions.
   /// Given a definition (instruction and definition index), this class
   /// follows the use-def chain to find successive suitable sources.
   /// The given source can be used to rewrite the definition into
   /// def = COPY src.
   ///
   /// For instance, let us consider the following snippet:
private:
     /// MachineRegisterInfo used to perform tracking.
     const MachineRegisterInfo &MRI;
     /// Optional TargetInstrInfo used to perform some complex
     /// tracking.
     const TargetInstrInfo *TII;
 
     /// \brief Dispatcher to the right underlying implementation of
     /// getNextSource.
-    bool getNextSourceImpl(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceImpl();
     /// \brief Specialized version of getNextSource for Copy instructions.
-    bool getNextSourceFromCopy(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceFromCopy();
     /// \brief Specialized version of getNextSource for Bitcast instructions.
-    bool getNextSourceFromBitcast(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceFromBitcast();
     /// \brief Specialized version of getNextSource for RegSequence
     /// instructions.
-    bool getNextSourceFromRegSequence(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceFromRegSequence();
     /// \brief Specialized version of getNextSource for InsertSubreg
     /// instructions.
-    bool getNextSourceFromInsertSubreg(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceFromInsertSubreg();
     /// \brief Specialized version of getNextSource for ExtractSubreg
     /// instructions.
-    bool getNextSourceFromExtractSubreg(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceFromExtractSubreg();
     /// \brief Specialized version of getNextSource for SubregToReg
     /// instructions.
-    bool getNextSourceFromSubregToReg(unsigned &SrcReg, unsigned &SrcSubReg);
+    ValueTrackerResult getNextSourceFromSubregToReg();
 
   public:
     /// \brief Create a ValueTracker instance for the value defined by \p Reg.
     /// \p DefSubReg represents the sub register index the value tracker will
     /// track. It does not need to match the sub register index used in the
     /// definition of \p Reg.
     /// \p UseAdvancedTracking specifies whether or not the value tracker looks
     /// through complex instructions. By default (false), it handles only copy
ValueTracker(const MachineInstr &MI, unsigned DefIdx, unsigned DefSubReg,
           UseAdvancedTracking(UseAdvancedTracking), MRI(MRI), TII(TII) {
       assert(DefIdx < Def->getDesc().getNumDefs() &&
              Def->getOperand(DefIdx).isReg() && "Invalid definition");
       Reg = Def->getOperand(DefIdx).getReg();
     }
 
     /// \brief Following the use-def chain, get the next available source
     /// for the tracked value.
-    /// When the returned value is not nullptr, \p SrcReg gives the register
-    /// that contain the tracked value.
-    /// \note The sub register index returned in \p SrcSubReg must be used
-    /// on \p SrcReg to access the actual value.
-    /// \return Unless the returned value is nullptr (i.e., no source found),
-    /// \p SrcReg gives the register of the next source used in the returned
-    /// instruction and \p SrcSubReg the sub-register index to be used on that
-    /// source to get the tracked value. When nullptr is returned, no
-    /// alternative source has been found.
-    const MachineInstr *getNextSource(unsigned &SrcReg, unsigned &SrcSubReg);
+    /// \return A ValueTrackerResult containing the a set of registers
+    /// and sub registers with tracked values. A ValueTrackerResult with
+    /// an empty set of registers means no source was found.
+    ValueTrackerResult getNextSource();
 
     /// \brief Get the last register where the initial value can be found.
     /// Initially this is the register of the definition.
     /// Then, after each successful call to getNextSource, this is the
     /// register of the last source.
     unsigned getReg() const { return Reg; }
   };
 }
bool PeepholeOptimizer::findNextSource(unsigned &Reg, unsigned &SubReg) {
   unsigned Src;
   unsigned SrcSubReg;
   bool ShouldRewrite = false;
 
   // Follow the chain of copies until we reach the top of the use-def chain
   // or find a more suitable source.
   ValueTracker ValTracker(Reg, DefSubReg, *MRI, !DisableAdvCopyOpt, TII);
   do {
-    unsigned CopySrcReg, CopySrcSubReg;
-    if (!ValTracker.getNextSource(CopySrcReg, CopySrcSubReg))
+    ValueTrackerResult Res = ValTracker.getNextSource();
+    if (!Res.isValid())
       break;
-    Src = CopySrcReg;
-    SrcSubReg = CopySrcSubReg;
+    Src = Res.getSrcReg(0);
+    SrcSubReg = Res.getSrcSubReg(0);
 
     // Do not extend the live-ranges of physical registers as they add
     // constraints to the register allocator.
     // Moreover, if we want to extend the live-range of a physical register,
     // unlike SSA virtual register, we will have to check that they are not
     // redefine before the related use.
     if (TargetRegisterInfo::isPhysicalRegister(Src))
       break;
virtual bool getNextRewritableSource(unsigned &SrcReg, unsigned &SrcSubReg,
     const MachineOperand &MODef = CopyLike.getOperand(0);
     TrackReg = MODef.getReg();
     TrackSubReg = MODef.getSubReg();
     return true;
   }
 
   /// \brief Rewrite the current source with \p NewReg and \p NewSubReg
   /// if possible.
-  /// \return True if the rewritting was possible, false otherwise.
+  /// \return True if the rewriting was possible, false otherwise.
   virtual bool RewriteCurrentSource(unsigned NewReg, unsigned NewSubReg) {
     if (!CopyLike.isCopy() || CurrentSrcIdx != 1)
       return false;
     MachineOperand &MOSrc = CopyLike.getOperand(CurrentSrcIdx);
     MOSrc.setReg(NewReg);
     MOSrc.setSubReg(NewSubReg);
     return true;
   }
+
+  /// \brief Rewrite the current source with \p NewSrcReg and \p NewSecSubReg
+  /// by creating a new COPY instruction. \p DefReg and \p DefSubReg contain the
+  /// definition to be rewritten from the original copylike instruction.
+  /// \return The new COPY if the rewriting was possible, nullptr otherwise.
+  /// This is needed to handle Uncoalescable copies, since they are copy
+  /// like instructions that aren't recognized by the register allocator.
+  virtual MachineInstr *RewriteCurrentSource(unsigned DefReg,
+                                             unsigned DefSubReg,
+                                             unsigned NewSrcReg,
+                                             unsigned NewSrcSubReg) {
+    return nullptr;
+  }
+};
+
+/// \brief Helper class to rewrite uncoalescable copy like instructions
+/// into new COPY (coalescable friendly) instructions.
+class UncoalescableRewriter : public CopyRewriter {
+protected:
+  const TargetInstrInfo &TII;
+  MachineRegisterInfo   &MRI;
+  /// The number of defs in the bitcast
+  unsigned NumDefs;
+
+public:
+  UncoalescableRewriter(MachineInstr &MI, const TargetInstrInfo &TII,
+                         MachineRegisterInfo &MRI)
+      : CopyRewriter(MI), TII(TII), MRI(MRI) {
+    NumDefs = MI.getDesc().getNumDefs();
+  }
+
+  /// \brief Get the next rewritable def source (TrackReg, TrackSubReg)
+  /// All such sources need to be considered rewritable in order to
+  /// rewrite a uncoalescable copy-like instruction. This method return
+  /// each definition that must be checked if rewritable.
+  ///
+  bool getNextRewritableSource(unsigned &SrcReg, unsigned &SrcSubReg,
+                               unsigned &TrackReg,
+                               unsigned &TrackSubReg) override {
+    // Find the next non-dead definition and continue from there.
+    if (CurrentSrcIdx == NumDefs)
+      return false;
+
+    while (CopyLike.getOperand(CurrentSrcIdx).isDead()) {
+      ++CurrentSrcIdx;
+      if (CurrentSrcIdx == NumDefs)
+        return false;
+    }
+
+    // What we track are the alternative sources of the definition.
+    const MachineOperand &MODef = CopyLike.getOperand(CurrentSrcIdx);
+    TrackReg = MODef.getReg();
+    TrackSubReg = MODef.getSubReg();
+
+    CurrentSrcIdx++;
+    return true;
+  }
+
+  /// \brief Rewrite the current source with \p NewSrcReg and \p NewSrcSubReg
+  /// by creating a new COPY instruction. \p DefReg and \p DefSubReg contain the
+  /// definition to be rewritten from the original copylike instruction.
+  /// \return The new COPY if the rewriting was possible, nullptr otherwise.
+  MachineInstr *RewriteCurrentSource(unsigned DefReg, unsigned DefSubReg,
+                                     unsigned NewSrcReg,
+                                     unsigned NewSrcSubReg) override {
+    assert(!TargetRegisterInfo::isPhysicalRegister(DefReg) &&
+           "We do not rewrite physical registers");
+
+    const TargetRegisterClass *DefRC = MRI.getRegClass(DefReg);
+    unsigned NewVR = MRI.createVirtualRegister(DefRC);
+
+    MachineInstr *NewCopy =
+        BuildMI(*CopyLike.getParent(), &CopyLike, CopyLike.getDebugLoc(),
+                TII.get(TargetOpcode::COPY), NewVR)
+            .addReg(NewSrcReg, 0, NewSrcSubReg);
+
+    NewCopy->getOperand(0).setSubReg(DefSubReg);
+    if (DefSubReg)
+      NewCopy->getOperand(0).setIsUndef();
+
+    MRI.replaceRegWith(DefReg, NewVR);
+    MRI.clearKillFlags(NewVR);
+
+    return NewCopy;
+  }
 };
 
 /// \brief Specialized rewriter for INSERT_SUBREG instruction.
 class InsertSubregRewriter : public CopyRewriter {
 public:
   InsertSubregRewriter(MachineInstr &MI) : CopyRewriter(MI) {
     assert(MI.isInsertSubreg() && "Invalid instruction");
   }
public:
   }
 };
 } // End namespace.
 
 /// \brief Get the appropriated CopyRewriter for \p MI.
 /// \return A pointer to a dynamically allocated CopyRewriter or nullptr
 /// if no rewriter works for \p MI.
 static CopyRewriter *getCopyRewriter(MachineInstr &MI,
-                                     const TargetInstrInfo &TII) {
+                                     const TargetInstrInfo &TII,
+                                     MachineRegisterInfo &MRI) {
+  // Handle uncoalescable copy-like instructions.
+  if (MI.isBitcast() || (MI.isRegSequenceLike() || MI.isInsertSubregLike() ||
+                         MI.isExtractSubregLike()))
+    return new UncoalescableRewriter(MI, TII, MRI);
+
   switch (MI.getOpcode()) {
   default:
     return nullptr;
   case TargetOpcode::COPY:
     return new CopyRewriter(MI);
   case TargetOpcode::INSERT_SUBREG:
     return new InsertSubregRewriter(MI);
   case TargetOpcode::EXTRACT_SUBREG:
assert(MI->getDesc().getNumDefs() == 1 &&
          "Coalescer can understand multiple defs?!");
   const MachineOperand &MODef = MI->getOperand(0);
   // Do not rewrite physical definitions.
   if (TargetRegisterInfo::isPhysicalRegister(MODef.getReg()))
     return false;
 
   bool Changed = false;
   // Get the right rewriter for the current copy.
-  std::unique_ptr<CopyRewriter> CpyRewriter(getCopyRewriter(*MI, *TII));
+  std::unique_ptr<CopyRewriter> CpyRewriter(getCopyRewriter(*MI, *TII, *MRI));
   // If none exists, bails out.
   if (!CpyRewriter)
     return false;
   // Rewrite each rewritable source.
   unsigned SrcReg, SrcSubReg, TrackReg, TrackSubReg;
   while (CpyRewriter->getNextRewritableSource(SrcReg, SrcSubReg, TrackReg,
                                               TrackSubReg)) {
     unsigned NewSrc = TrackReg;
     unsigned NewSubReg = TrackSubReg;
-    // Try to find a more suitable source.
-    // If we failed to do so, or get the actual source,
-    // move to the next source.
+    // Try to find a more suitable source. If we failed to do so, or get the
+    // actual source, move to the next source.
     if (!findNextSource(NewSrc, NewSubReg) || SrcReg == NewSrc)
       continue;
     // Rewrite source.
     if (CpyRewriter->RewriteCurrentSource(NewSrc, NewSubReg)) {
       // We may have extended the live-range of NewSrc, account for that.
       MRI->clearKillFlags(NewSrc);
       Changed = true;
     }

 bool PeepholeOptimizer::optimizeUncoalescableCopy(
     MachineInstr *MI, SmallPtrSetImpl<MachineInstr *> &LocalMIs) {
   assert(MI && isUncoalescableCopy(*MI) && "Invalid argument");
 
   // Check if we can rewrite all the values defined by this instruction.
   SmallVector<
       std::pair<TargetInstrInfo::RegSubRegPair, TargetInstrInfo::RegSubRegPair>,
       4> RewritePairs;
-  for (const MachineOperand &MODef : MI->defs()) {
-    if (MODef.isDead())
-      // We can ignore those.
-      continue;
+  // Get the right rewriter for the current copy.
+  std::unique_ptr<CopyRewriter> CpyRewriter(getCopyRewriter(*MI, *TII, *MRI));
+  // If none exists, bails out.
+  if (!CpyRewriter)
+    return false;
 
+  // Rewrite each rewritable source by generating new COPYs. This works
+  // differently from optimizeCoalescableCopy since it first makes sure that all
+  // definitions can be rewritten.
+  unsigned SrcReg, SrcSubReg, TrackReg, TrackSubReg;
+  while (CpyRewriter->getNextRewritableSource(SrcReg, SrcSubReg, TrackReg,
+                                              TrackSubReg)) {
     // If a physical register is here, this is probably for a good reason.
     // Do not rewrite that.
-    if (TargetRegisterInfo::isPhysicalRegister(MODef.getReg()))
+    if (TargetRegisterInfo::isPhysicalRegister(TrackReg))
       return false;
 
     // If we do not know how to rewrite this definition, there is no point
     // in trying to kill this instruction.
-    TargetInstrInfo::RegSubRegPair Def(MODef.getReg(), MODef.getSubReg());
+    TargetInstrInfo::RegSubRegPair Def(TrackReg, TrackSubReg);
     TargetInstrInfo::RegSubRegPair Src = Def;
     if (!findNextSource(Src.Reg, Src.SubReg))
       return false;
+
     RewritePairs.push_back(std::make_pair(Def, Src));
   }
+
   // The change is possible for all defs, do it.
   for (const auto &PairDefSrc : RewritePairs) {
     const auto &Def = PairDefSrc.first;
     const auto &Src = PairDefSrc.second;
+
     // Rewrite the "copy" in a way the register coalescer understands.
-    assert(!TargetRegisterInfo::isPhysicalRegister(Def.Reg) &&
-           "We do not rewrite physical registers");
-    const TargetRegisterClass *DefRC = MRI->getRegClass(Def.Reg);
-    unsigned NewVR = MRI->createVirtualRegister(DefRC);
-    MachineInstr *NewCopy = BuildMI(*MI->getParent(), MI, MI->getDebugLoc(),
-                                    TII->get(TargetOpcode::COPY),
-                                    NewVR).addReg(Src.Reg, 0, Src.SubReg);
-    NewCopy->getOperand(0).setSubReg(Def.SubReg);
-    if (Def.SubReg)
-      NewCopy->getOperand(0).setIsUndef();
-    LocalMIs.insert(NewCopy);
-    MRI->replaceRegWith(Def.Reg, NewVR);
-    MRI->clearKillFlags(NewVR);
-    // We extended the lifetime of Src.
-    // Clear the kill flags to account for that.
+    MachineInstr *NewCopy = CpyRewriter->RewriteCurrentSource(
+        Def.Reg, Def.SubReg, Src.Reg, Src.SubReg);
+    assert(NewCopy && "Should be able to always generate a new copy");
+
+    // We extended the lifetime of Src and clear the kill flags to
+    // account for that.
     MRI->clearKillFlags(Src.Reg);
+    LocalMIs.insert(NewCopy);
   }
   // MI is now dead.
   MI->eraseFromParent();
   ++NumUncoalescableCopies;
   return true;
 }
 
 /// isLoadFoldable - Check whether MI is a candidate for folding into a later
for (MachineBasicBlock::iterator
         }
       }
     }
   }
 
   return Changed;
 }
 
-bool ValueTracker::getNextSourceFromCopy(unsigned &SrcReg,
-                                         unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceFromCopy() {
   assert(Def->isCopy() && "Invalid definition");
   // Copy instruction are supposed to be: Def = Src.
   // If someone breaks this assumption, bad things will happen everywhere.
   assert(Def->getNumOperands() == 2 && "Invalid number of operands");
 
   if (Def->getOperand(DefIdx).getSubReg() != DefSubReg)
     // If we look for a different subreg, it means we want a subreg of src.
     // Bails as we do not support composing subreg yet.
-    return false;
+    return ValueTrackerResult();
   // Otherwise, we want the whole source.
   const MachineOperand &Src = Def->getOperand(1);
-  SrcReg = Src.getReg();
-  SrcSubReg = Src.getSubReg();
-  return true;
+  return ValueTrackerResult(Src.getReg(), Src.getSubReg());
 }
 
-bool ValueTracker::getNextSourceFromBitcast(unsigned &SrcReg,
-                                            unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceFromBitcast() {
   assert(Def->isBitcast() && "Invalid definition");
 
   // Bail if there are effects that a plain copy will not expose.
   if (Def->hasUnmodeledSideEffects())
-    return false;
+    return ValueTrackerResult();
 
   // Bitcasts with more than one def are not supported.
   if (Def->getDesc().getNumDefs() != 1)
-    return false;
+    return ValueTrackerResult();
   if (Def->getOperand(DefIdx).getSubReg() != DefSubReg)
     // If we look for a different subreg, it means we want a subreg of the src.
     // Bails as we do not support composing subreg yet.
-    return false;
+    return ValueTrackerResult();
 
   unsigned SrcIdx = Def->getNumOperands();
   for (unsigned OpIdx = DefIdx + 1, EndOpIdx = SrcIdx; OpIdx != EndOpIdx;
        ++OpIdx) {
     const MachineOperand &MO = Def->getOperand(OpIdx);
     if (!MO.isReg() || !MO.getReg())
       continue;
     assert(!MO.isDef() && "We should have skipped all the definitions by now");
     if (SrcIdx != EndOpIdx)
       // Multiple sources?
-      return false;
+      return ValueTrackerResult();
     SrcIdx = OpIdx;
   }
   const MachineOperand &Src = Def->getOperand(SrcIdx);
-  SrcReg = Src.getReg();
-  SrcSubReg = Src.getSubReg();
-  return true;
+  return ValueTrackerResult(Src.getReg(), Src.getSubReg());
 }
 
-bool ValueTracker::getNextSourceFromRegSequence(unsigned &SrcReg,
-                                                unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceFromRegSequence() {
   assert((Def->isRegSequence() || Def->isRegSequenceLike()) &&
          "Invalid definition");
 
   if (Def->getOperand(DefIdx).getSubReg())
     // If we are composing subreg, bails out.
     // The case we are checking is Def.<subreg> = REG_SEQUENCE.
     // This should almost never happen as the SSA property is tracked at
     // the register level (as opposed to the subreg level).
     // I.e.,
     // Def.sub0 =
     // Def.sub1 =
     // is a valid SSA representation for Def.sub0 and Def.sub1, but not for
     // Def. Thus, it must not be generated.
     // However, some code could theoretically generates a single
     // Def.sub0 (i.e, not defining the other subregs) and we would
     // have this case.
     // If we can ascertain (or force) that this never happens, we could
     // turn that into an assertion.
-    return false;
+    return ValueTrackerResult();
 
   if (!TII)
     // We could handle the REG_SEQUENCE here, but we do not want to
     // duplicate the code from the generic TII.
-    return false;
+    return ValueTrackerResult();
 
   SmallVector<TargetInstrInfo::RegSubRegPairAndIdx, 8> RegSeqInputRegs;
   if (!TII->getRegSequenceInputs(*Def, DefIdx, RegSeqInputRegs))
-    return false;
+    return ValueTrackerResult();
 
   // We are looking at:
   // Def = REG_SEQUENCE v0, sub0, v1, sub1, ...
   // Check if one of the operand defines the subreg we are interested in.
   for (auto &RegSeqInput : RegSeqInputRegs) {
     if (RegSeqInput.SubIdx == DefSubReg) {
       if (RegSeqInput.SubReg)
         // Bails if we have to compose sub registers.
-        return false;
+        return ValueTrackerResult();
 
-      SrcReg = RegSeqInput.Reg;
-      SrcSubReg = RegSeqInput.SubReg;
-      return true;
+      return ValueTrackerResult(RegSeqInput.Reg, RegSeqInput.SubReg);
     }
   }
 
   // If the subreg we are tracking is super-defined by another subreg,
   // we could follow this value. However, this would require to compose
   // the subreg and we do not do that for now.
-  return false;
+  return ValueTrackerResult();
 }
 
-bool ValueTracker::getNextSourceFromInsertSubreg(unsigned &SrcReg,
-                                                 unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceFromInsertSubreg() {
   assert((Def->isInsertSubreg() || Def->isInsertSubregLike()) &&
          "Invalid definition");
 
   if (Def->getOperand(DefIdx).getSubReg())
     // If we are composing subreg, bails out.
     // Same remark as getNextSourceFromRegSequence.
     // I.e., this may be turned into an assert.
-    return false;
+    return ValueTrackerResult();
 
   if (!TII)
     // We could handle the REG_SEQUENCE here, but we do not want to
     // duplicate the code from the generic TII.
-    return false;
+    return ValueTrackerResult();
 
   TargetInstrInfo::RegSubRegPair BaseReg;
   TargetInstrInfo::RegSubRegPairAndIdx InsertedReg;
   if (!TII->getInsertSubregInputs(*Def, DefIdx, BaseReg, InsertedReg))
-    return false;
+    return ValueTrackerResult();
 
   // We are looking at:
   // Def = INSERT_SUBREG v0, v1, sub1
   // There are two cases:
   // 1. DefSubReg == sub1, get v1.
   // 2. DefSubReg != sub1, the value may be available through v0.
 
   // #1 Check if the inserted register matches the required sub index.
   if (InsertedReg.SubIdx == DefSubReg) {
-    SrcReg = InsertedReg.Reg;
-    SrcSubReg = InsertedReg.SubReg;
-    return true;
+    return ValueTrackerResult(InsertedReg.Reg, InsertedReg.SubReg);
   }
   // #2 Otherwise, if the sub register we are looking for is not partial
   // defined by the inserted element, we can look through the main
   // register (v0).
   const MachineOperand &MODef = Def->getOperand(DefIdx);
   // If the result register (Def) and the base register (v0) do not
   // have the same register class or if we have to compose
   // subregisters, bails out.
   if (MRI.getRegClass(MODef.getReg()) != MRI.getRegClass(BaseReg.Reg) ||
       BaseReg.SubReg)
-    return false;
+    return ValueTrackerResult();
 
   // Get the TRI and check if the inserted sub-register overlaps with the
   // sub-register we are tracking.
   const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
   if (!TRI ||
       (TRI->getSubRegIndexLaneMask(DefSubReg) &
        TRI->getSubRegIndexLaneMask(InsertedReg.SubIdx)) != 0)
-    return false;
+    return ValueTrackerResult();
   // At this point, the value is available in v0 via the same subreg
   // we used for Def.
-  SrcReg = BaseReg.Reg;
-  SrcSubReg = DefSubReg;
-  return true;
+  return ValueTrackerResult(BaseReg.Reg, DefSubReg);
 }
 
-bool ValueTracker::getNextSourceFromExtractSubreg(unsigned &SrcReg,
-                                                  unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceFromExtractSubreg() {
   assert((Def->isExtractSubreg() ||
           Def->isExtractSubregLike()) && "Invalid definition");
   // We are looking at:
   // Def = EXTRACT_SUBREG v0, sub0
 
   // Bails if we have to compose sub registers.
   // Indeed, if DefSubReg != 0, we would have to compose it with sub0.
   if (DefSubReg)
-    return false;
+    return ValueTrackerResult();
 
   if (!TII)
     // We could handle the EXTRACT_SUBREG here, but we do not want to
     // duplicate the code from the generic TII.
-    return false;
+    return ValueTrackerResult();
 
   TargetInstrInfo::RegSubRegPairAndIdx ExtractSubregInputReg;
   if (!TII->getExtractSubregInputs(*Def, DefIdx, ExtractSubregInputReg))
-    return false;
+    return ValueTrackerResult();
 
   // Bails if we have to compose sub registers.
   // Likewise, if v0.subreg != 0, we would have to compose v0.subreg with sub0.
   if (ExtractSubregInputReg.SubReg)
-    return false;
+    return ValueTrackerResult();
   // Otherwise, the value is available in the v0.sub0.
-  SrcReg = ExtractSubregInputReg.Reg;
-  SrcSubReg = ExtractSubregInputReg.SubIdx;
-  return true;
+  return ValueTrackerResult(ExtractSubregInputReg.Reg, ExtractSubregInputReg.SubIdx);
 }
 
-bool ValueTracker::getNextSourceFromSubregToReg(unsigned &SrcReg,
-                                                unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceFromSubregToReg() {
   assert(Def->isSubregToReg() && "Invalid definition");
   // We are looking at:
   // Def = SUBREG_TO_REG Imm, v0, sub0
 
   // Bails if we have to compose sub registers.
   // If DefSubReg != sub0, we would have to check that all the bits
   // we track are included in sub0 and if yes, we would have to
   // determine the right subreg in v0.
   if (DefSubReg != Def->getOperand(3).getImm())
-    return false;
+    return ValueTrackerResult();
   // Bails if we have to compose sub registers.
   // Likewise, if v0.subreg != 0, we would have to compose it with sub0.
   if (Def->getOperand(2).getSubReg())
-    return false;
+    return ValueTrackerResult();
 
-  SrcReg = Def->getOperand(2).getReg();
-  SrcSubReg = Def->getOperand(3).getImm();
-  return true;
+  return ValueTrackerResult(Def->getOperand(2).getReg(),
+                            Def->getOperand(3).getImm());
 }
 
-bool ValueTracker::getNextSourceImpl(unsigned &SrcReg, unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSourceImpl() {
   assert(Def && "This method needs a valid definition");
 
   assert(
       (DefIdx < Def->getDesc().getNumDefs() || Def->getDesc().isVariadic()) &&
       Def->getOperand(DefIdx).isDef() && "Invalid DefIdx");
   if (Def->isCopy())
-    return getNextSourceFromCopy(SrcReg, SrcSubReg);
+    return getNextSourceFromCopy();
   if (Def->isBitcast())
-    return getNextSourceFromBitcast(SrcReg, SrcSubReg);
+    return getNextSourceFromBitcast();
   // All the remaining cases involve "complex" instructions.
   // Bails if we did not ask for the advanced tracking.
   if (!UseAdvancedTracking)
-    return false;
+    return ValueTrackerResult();
   if (Def->isRegSequence() || Def->isRegSequenceLike())
-    return getNextSourceFromRegSequence(SrcReg, SrcSubReg);
+    return getNextSourceFromRegSequence();
   if (Def->isInsertSubreg() || Def->isInsertSubregLike())
-    return getNextSourceFromInsertSubreg(SrcReg, SrcSubReg);
+    return getNextSourceFromInsertSubreg();
   if (Def->isExtractSubreg() || Def->isExtractSubregLike())
-    return getNextSourceFromExtractSubreg(SrcReg, SrcSubReg);
+    return getNextSourceFromExtractSubreg();
   if (Def->isSubregToReg())
-    return getNextSourceFromSubregToReg(SrcReg, SrcSubReg);
-  return false;
+    return getNextSourceFromSubregToReg();
+  return ValueTrackerResult();
 }
 
-const MachineInstr *ValueTracker::getNextSource(unsigned &SrcReg,
-                                                unsigned &SrcSubReg) {
+ValueTrackerResult ValueTracker::getNextSource() {
   // If we reach a point where we cannot move up in the use-def chain,
   // there is nothing we can get.
   if (!Def)
-    return nullptr;
+    return ValueTrackerResult();
 
-  const MachineInstr *PrevDef = nullptr;
-  // Try to find the next source.
-  if (getNextSourceImpl(SrcReg, SrcSubReg)) {
+  ValueTrackerResult Res = getNextSourceImpl();
+  if (Res.isValid()) {
     // Update definition, definition index, and subregister for the
     // next call of getNextSource.
     // Update the current register.
-    Reg = SrcReg;
-    // Update the return value before moving up in the use-def chain.
-    PrevDef = Def;
+    bool OneRegSrc = Res.getNumSources() == 1;
+    if (OneRegSrc)
+      Reg = Res.getSrcReg(0);
+    // Update the result before moving up in the use-def chain
+    // with the instruction containing the last found sources.
+    Res.setInst(Def);
+
     // If we can still move up in the use-def chain, move to the next
     // defintion.
-    if (!TargetRegisterInfo::isPhysicalRegister(Reg)) {
+    if (!TargetRegisterInfo::isPhysicalRegister(Reg) && OneRegSrc) {
       Def = MRI.getVRegDef(Reg);
       DefIdx = MRI.def_begin(Reg).getOperandNo();
-      DefSubReg = SrcSubReg;
-      return PrevDef;
+      DefSubReg = Res.getSrcSubReg(0);
+      return Res;
     }
   }
   // If we end up here, this means we will not be able to find another source
-  // for the next iteration.
-  // Make sure any new call to getNextSource bails out early by cutting the
-  // use-def chain.
+  // for the next iteration. Make sure any new call to getNextSource bails out
+  // early by cutting the use-def chain.
   Def = nullptr;
-  return PrevDef;
+  return Res;
 }