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

[MaskRay]

Summary

Reland 7a8b0b1595e7dc878b48cf9bbaa652087a6895db, with a fix that checks
!E.value().empty() to avoid inserting a zero to SlotRemap.

Debugged by rnk@ in https://bugs.chromium.org/p/chromium/issues/detail?id=1045650#c33 . I don't know how to get a WinEH reproduce.

[MaskRay]

test/CodeGen/X86/catchpad-lifetime.ll (added in rL257158 @majnemer) appears to be bit-rotted. If I delete:

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];

Not test will break.

[merge_guards_bot]

Unit tests: pass. 62248 tests passed, 0 failed and 816 were skipped.

clang-tidy: pass.

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-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rnk]

In D73510#1843302, @MaskRay wrote:

test/CodeGen/X86/catchpad-lifetime.ll (added in rL257158 @majnemer) appears to be bit-rotted. If I delete:
...
Not test will break.

I think it used to be covered by x64 tests, but then not long after in rG1ef654024f490b403494265c968570ca6bd800f3 we made most x64 catch objects fixed stack objects, which stack coloring should ignore. I ported one test to x86 in rGc7feb6b36aa8dac3cd20ffb9ad5980693ea9916e, but I'm not sure that one actually exercises this code.

[rnk]

lgtm

llvm/lib/CodeGen/StackColoring.cpp

1081

Are you sure it wouldn't be better to do *SlotRemap.find(E.index())? Then it is clear that the operation is readonly. As is, I agree that it is impossible for this to insert a value, the vector for this index will only be non-empty if there is already an entry in the table.

[MaskRay]

Replace operator[] with

MF->getPSVManager().getFixedStack(SlotRemap.find(E.index())->second);

[merge_guards_bot]

Unit tests: pass. 62248 tests passed, 0 failed and 816 were skipped.

clang-tidy: pass.

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-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

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))
+    if (!E.value().empty()) {
+      const PseudoSourceValue *NewSV =
+          MF->getPSVManager().getFixedStack(SlotRemap.find(E.index())->second);
+      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