[StackColoring] Remap PseudoSourceValue frame indices via MachineFunction::getPSVManager()

[merge_guards_bot]

Unit tests: pass. 62027 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[dantrushin]

LGTM

llvm/include/llvm/CodeGen/PseudoSourceValue.h

public:
   /// Return true if the memory pointed to by this PseudoSourceValue can ever
   /// alias an LLVM IR Value.
   virtual bool mayAlias(const MachineFrameInfo *) const;
 };
 
 /// A specialized PseudoSourceValue for holding FixedStack values, which must
 /// include a frame index.
 class FixedStackPseudoSourceValue : public PseudoSourceValue {
-  int FI;
+  const int FI;
 
 public:
   explicit FixedStackPseudoSourceValue(int FI, const TargetInstrInfo &TII)
       : PseudoSourceValue(FixedStack, TII), FI(FI) {}
 
   static bool classof(const PseudoSourceValue *V) {
     return V->kind() == FixedStack;
   }
 
   bool isConstant(const MachineFrameInfo *MFI) const override;
 
   bool isAliased(const MachineFrameInfo *MFI) const override;
 
   bool mayAlias(const MachineFrameInfo *) const override;
 
   void printCustom(raw_ostream &OS) const override;
 
   int getFrameIndex() const { return FI; }
-  void setFrameIndex(int FI) { this->FI = FI; }
 };
 
 class CallEntryPseudoSourceValue : public PseudoSourceValue {
 protected:
   CallEntryPseudoSourceValue(unsigned Kind, const TargetInstrInfo &TII);
 
 public:
   bool isConstant(const MachineFrameInfo *) const override;

llvm/lib/CodeGen/StackColoring.cpp

for (const std::pair<int, int> &SI : SlotRemap) {
 
     // Note that this will not replace uses in MMOs (which we'll update below),
     // or anywhere else (which is why we won't delete the original
     // instruction).
     FromAI->replaceAllUsesWith(Inst);
   }
 
   // Remap all instructions to the new stack slots.
+  std::vector<std::vector<MachineMemOperand *>> SSRefs(MFI->getObjectIndexEnd());
   for (MachineBasicBlock &BB : *MF)
     for (MachineInstr &I : BB) {
       // Skip lifetime markers. We'll remove them soon.
       if (I.getOpcode() == TargetOpcode::LIFETIME_START ||
           I.getOpcode() == TargetOpcode::LIFETIME_END)
         continue;
 
       // Update the MachineMemOperand to use the new alloca.
#endif
         MO.setIndex(ToSlot);
         FixedInstr++;
       }
 
       // We adjust AliasAnalysis information for merged stack slots.
       SmallVector<MachineMemOperand *, 2> NewMMOs;
       bool ReplaceMemOps = false;
       for (MachineMemOperand *MMO : I.memoperands()) {
+        // Collect MachineMemOperands which reference
+        // FixedStackPseudoSourceValues with old frame indices.
         if (const auto *FSV = dyn_cast_or_null<FixedStackPseudoSourceValue>(
                 MMO->getPseudoValue())) {
           int FI = FSV->getFrameIndex();
           auto To = SlotRemap.find(FI);
           if (To != SlotRemap.end())
-            const_cast<FixedStackPseudoSourceValue *>(FSV)->setFrameIndex(
-                To->second);
+            SSRefs[FI].push_back(MMO);
         }
 
         // If this memory location can be a slot remapped here,
         // we remove AA information.
         bool MayHaveConflictingAAMD = false;
         if (MMO->getAAInfo()) {
           if (const Value *MMOV = MMO->getValue()) {
             SmallVector<Value *, 4> Objs;
#endif
       }
 
       // If any memory operand is updated, set memory references of
       // this instruction.
       if (ReplaceMemOps)
         I.setMemRefs(*MF, NewMMOs);
     }
 
+  // Rewrite MachineMemOperands that reference old frame indices.
+  for (auto E : enumerate(SSRefs)) {
+    const PseudoSourceValue *NewSV =
+        MF->getPSVManager().getFixedStack(SlotRemap[E.index()]);
+    for (MachineMemOperand *Ref : E.value())
+      Ref->setValue(NewSV);
+  }
+
   // Update the location of C++ catch objects for the MSVC personality routine.
   if (WinEHFuncInfo *EHInfo = MF->getWinEHFuncInfo())
     for (WinEHTryBlockMapEntry &TBME : EHInfo->TryBlockMap)
       for (WinEHHandlerType &H : TBME.HandlerArray)
         if (H.CatchObj.FrameIndex != std::numeric_limits<int>::max() &&
             SlotRemap.count(H.CatchObj.FrameIndex))
           H.CatchObj.FrameIndex = SlotRemap[H.CatchObj.FrameIndex];
 

llvm/test/CodeGen/PowerPC/stack-coloring-vararg.mir

-# RUN: llc -o - %s -start-before=stack-coloring
+# RUN: llc -run-pass=stack-coloring %s -o - | FileCheck %s
+
+## Test %stack.1 is merged into %stack.0 and there is no MemoryMemOperand
+## referencing %stack.1. This regression test is sensitive to the StackColoring
+## algorithm. Please adjust or delete this test if the merging strategy
+## changes.
+
+# CHECK:    {{^}}stack:
+# CHECK-NEXT: - { id: 0,
+# CHECK-NOT:  - { id: 1,
+# CHECK:      - { id: 2,
+# CHECK-NOT: %stack.1
+
 --- |
   ; ModuleID = '<stdin>'
   source_filename = "<stdin>"
   target datalayout = "E-m:e-p:32:32-i64:64-n32"
   target triple = "powerpc-unknown-freebsd13.0"
 
   %struct.__va_list_tag = type { i8, i8, i16, i8*, i8* }
   ; Function Attrs: argmemonly nounwind willreturn