Factor PrologEpilogInserter around spilling, frame finalization, and scavenging

[dschuff]

Summary

PrologEpilogInserter has these 3 phases, which are related, but not
all of them are needed by all targets. This patch reorganizes PEI's
varous functions around those phases for more clear separation, and adds
more useful assertions based on which parts must support virtual
registers, callee-saved register, and scavenging.

[qcolombet]

Hi Derek,

Thanks for keep pushing on that.

The HasVirtualRegister will not change for a given target, thus, I would rather have it checked only at the instantiation of the class.
In other words, instead of having two if blocks in runOnMachineFunction, I would suggest calling a helper object that would have been instantiated differently according to the mode.

E.g., instead of doing
for each function
if (HasVirtual) {

// handlingA

} else {

// handlingAbis

}

I would do:
Once for the Module:
if (HasVirtual) {
helper = helperA
} else {
helper = helperAbis // where helperAbis is a sub class of helperA and helperA does not nothing by default.
}

for each function

helper->callA

What do you think?

Cheers,
-Quentin

lib/CodeGen/PrologEpilogInserter.cpp

189

HasVirtualRegisters should be an argument of the pass and we should have an assert HasVirtualRegisters or ! Fn.getRegInfo().getNumVirtRegs().
The rationale is that some backends might end up having virtual registers here whereas they should not.

[dschuff]

The problem with checking at pass instantiation time is that (I think) you can't query the target. It's only the MachineFunction object that has the knowledge of what target you are compiling for, and you don't get that until runOnMachineFunction (or runOnFunction in the base MachineFunctionPass).

lib/CodeGen/PrologEpilogInserter.cpp

189

See the bottom comment about arguments to the pass; here I'll just note that the asserts on lines 161 and 198 do enforce the invariants we care about (assuming there are no virtual targets that require virtual registers but also have CSRs).

[dschuff]

• Configure PEI at instantiation time and factor CSR out of pass class

[dschuff]

Here's a stab in that direction (just for CSR spilling, we'd do the same kind of thing for scavenging before actually landing it).
There are several changes:

• Pass instantiation gets a TargetMachine so it can decide whether the target wants virtual registers. We'd need to add some kind of target hook for that.
• SpillCalleeSavedRegisters becomes a function pointer that does nothing, or the usual thing. This version of the patch makes it a non-member function and passes the data it needs explicitly (before it was just setting member variables). I'm on the fence about whether that's better; I like that it's easier to see the dependencies, but it's 5 arguments to pass around.

[dschuff]

Ping!
(can't wait to delete my copy of PrologEpilogInserter from the wasm backend :)

[qcolombet]

lib/CodeGen/PrologEpilogInserter.cpp

52

I’d say for consistency with LLVM codebase we should stick to typedef.

69

Indeed :).

228

I still don’t like this part of the patch.
We could do another helper function here, couldn’t we?

663

Unrelated change.

[dschuff]

lib/CodeGen/PrologEpilogInserter.cpp

228

Yes, sorry if I wasn't clear enough in my comment from 3/29; the intent was to get your opinion on the general direction using CSR spilling and then apply the same idea to scavenging. Since you didn't have any objection to the rest of it, I'll go ahead and do that now :)

[dschuff]

• Configure PEI at instantiation time and factor CSR out of pass class
• Also factor out scavenging, fix pass instantiation

[dschuff]

• remove HasVirtualRegisters and its assert

[dschuff]

OK, Here's PEI with register scavenging also factored out.

1. It's still missing the target hook. I'm looking into that.
2. Requiring the query of the TargetMachine at pass instantiation time makes the logic in TargetPassConfig a bit ugly because passes are substituted or disabled by ID but now we have instantiate the pass in TargetPassConfig directly. Maybe there's a better way but this way at least preserves the ability of targets to disable, override, or substitute PrologEpilogInserterID.

lib/CodeGen/PrologEpilogInserter.cpp

663

It's actually not an unrelated change. MinCSFrameIndex and MaxCSFrameIndex are unsigned. Previously MinCSFrameIndex was initialized to INT_MAX, and when I moved that i made it std::numeric_limits<unsigned>::max() to match the type. But then this casting is wrong because MinCSFI becomes -1 instead of a big number when there are no CSRs. So I decided to fix this type mismatch rather than keep the initializer type mismatch just for the purpose of also being able to keep this type mismatch.

[dschuff]

lib/CodeGen/PrologEpilogInserter.cpp

69

For the target hook you could almost use TargetRegisterInfo::getCalleeSavedRegs() (if CSRs is empty, then you can skip spilling) but that requires a MachineFunction because it depends on the calling convention. We could add something like TargetRegisterInfo::usesCalleeSavedRegs() or even come full circle to something like in D15394 :)

[qcolombet]

lib/CodeGen/PrologEpilogInserter.cpp

69

I don’t like the name in D15394, in particular the mention of register allocation (it is antinomic to me to have both virtual register and after regalloc in the same name :)).
The problem with usesCalleeSavedRegs is that it does not convey the information that we want or do not want to scavenge the virtual registers used for frame.

87

I do not like that UsesCalleeSaves implies all vregs are allocated. That is the same naming problem I mentioned previously, We need to come up with a name for the hook that conveys all that information at once. usesPhysReg for PEI?

[dschuff]

lib/CodeGen/PrologEpilogInserter.cpp

87

Do you mean usesPhysRegForPEI? or we could be more generic and do usesVirtRegsOnly?.

One issue with the code as-is, is that there's no way to get a targetRegisterInfo without a function (because it hangs off the subtarget, which can be different for each function). We could add a new method to TargetMachine itself but that seems like a big hammer. I'm not sure what else you could do though if you really want to do the pass configuration at instantiation time like this.

[qcolombet]

lib/CodeGen/PrologEpilogInserter.cpp

87

Do you mean usesPhysRegForPEI?

Yes, that’s what I meant, sorry.

I don’t think this should be in TargetRegisterInfo. I believe this property should hold for the target and not differ between sub targets, otherwise we are mixing some strange code :).
A hook in TargetMachine maybe?

[dschuff]

• Use TargetMachine hook

[dschuff]

OK; I think this might actually be ready to go then...

[qcolombet]

Hi Derek,

Thanks for pushing this all the way through!

LGTM.

Cheers,
-Quentin

[dschuff]

Thanks for all the reviews!

include/llvm/CodeGen/Passes.h

public:
   void disablePass(AnalysisID PassID) {
     substitutePass(PassID, IdentifyingPassPtr());
   }
 
   /// Return the pass substituted for StandardID by the target.
   /// If no substitution exists, return StandardID.
   IdentifyingPassPtr getPassSubstitution(AnalysisID StandardID) const;
 
+  /// Return true if the pass has been substituted by the target or
+  /// overridden on the command line.
+  bool isPassSubstitutedOrOverridden(AnalysisID ID) const;
+
   /// Return true if the optimized regalloc pipeline is enabled.
   bool getOptimizeRegAlloc() const;
 
   /// Return true if shrink wrapping is enabled.
   bool getEnableShrinkWrap() const;
 
   /// Return true if the default global register allocator is in use and
   /// has not be overriden on the command line with '-regalloc=...'
/// MachineDominanaceFrontier - This pass is a machine dominators analysis pass.
   /// PBQPRegisterAllocation Pass - This pass implements the Partitioned Boolean
   /// Quadratic Prograaming (PBQP) based register allocator.
   ///
   FunctionPass *createDefaultPBQPRegisterAllocator();
 
   /// PrologEpilogCodeInserter - This pass inserts prolog and epilog code,
   /// and eliminates abstract frame references.
   extern char &PrologEpilogCodeInserterID;
+  MachineFunctionPass *createPrologEpilogInserterPass(const TargetMachine *TM);
 
   /// ExpandPostRAPseudos - This pass expands pseudo instructions after
   /// register allocation.
   extern char &ExpandPostRAPseudosID;
 
   /// createPostRAScheduler - This pass performs post register allocation
   /// scheduling.
   extern char &PostRASchedulerID;

lib/CodeGen/Passes.cpp

 IdentifyingPassPtr TargetPassConfig::getPassSubstitution(AnalysisID ID) const {
   DenseMap<AnalysisID, IdentifyingPassPtr>::const_iterator
     I = Impl->TargetPasses.find(ID);
   if (I == Impl->TargetPasses.end())
     return ID;
   return I->second;
 }
 
+bool TargetPassConfig::isPassSubstitutedOrOverridden(AnalysisID ID) const {
+  IdentifyingPassPtr TargetID = getPassSubstitution(ID);
+  IdentifyingPassPtr FinalPtr = overridePass(ID, TargetID);
+  return !FinalPtr.isValid() || FinalPtr.isInstance() ||
+      FinalPtr.getID() != ID;
+}
+
 /// Add a pass to the PassManager if that pass is supposed to be run.  If the
 /// Started/Stopped flags indicate either that the compilation should start at
 /// a later pass or that it should stop after an earlier pass, then do not add
 /// the pass.  Finally, compare the current pass against the StartAfter
 /// and StopAfter options and change the Started/Stopped flags accordingly.
 void TargetPassConfig::addPass(Pass *P, bool verifyAfter, bool printAfter) {
   assert(!Initialized && "PassConfig is immutable");
 
void TargetPassConfig::addMachinePasses() {
 
   // Run post-ra passes.
   addPostRegAlloc();
 
   // Insert prolog/epilog code.  Eliminate abstract frame index references...
   if (getOptLevel() != CodeGenOpt::None)
     addPass(&ShrinkWrapID);
 
-  addPass(&PrologEpilogCodeInserterID);
+  // Prolog/Epilog inserter needs a TargetMachine to instantiate. But only
+  // do so if it hasn't been disabled, substituted, or overridden.
+  if (!isPassSubstitutedOrOverridden(&PrologEpilogCodeInserterID))
+      addPass(createPrologEpilogInserterPass(TM));
 
   /// Add passes that optimize machine instructions after register allocation.
   if (getOptLevel() != CodeGenOpt::None)
     addMachineLateOptimization();
 
   // Expand pseudo instructions before second scheduling pass.
   addPass(&ExpandPostRAPseudosID);
 

lib/CodeGen/PrologEpilogInserter.cpp

 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include <climits>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "pei"
 
+typedef SmallVector<MachineBasicBlock *, 4> MBBVector;
+static void doSpillCalleeSavedRegs(MachineFunction &MF, RegScavenger *RS,
+                                   unsigned &MinCSFrameIndex,
+                                   unsigned &MaxCXFrameIndex,
+                                   const MBBVector &SaveBlocks,
+                                   const MBBVector &RestoreBlocks);
+
+static void doScavengeFrameVirtualRegs(MachineFunction &MF, RegScavenger *RS);
+
 namespace {
 class PEI : public MachineFunctionPass {
 public:
   static char ID;
-  PEI() : MachineFunctionPass(ID) {
+  explicit PEI(const TargetMachine *TM = nullptr) : MachineFunctionPass(ID) {
     initializePEIPass(*PassRegistry::getPassRegistry());
+
+    // XXX replace this with some sort of target hook
+    if (TM && (TM->getTargetTriple().getArch() == Triple::wasm32 ||
+               TM->getTargetTriple().getArch() == Triple::wasm64)) {
+      SpillCalleeSavedRegisters = [](MachineFunction &, RegScavenger *,
+                                     unsigned &, unsigned &, const MBBVector &,
+                                     const MBBVector &) {};
+      ScavengeFrameVirtualRegs = [](MachineFunction &, RegScavenger *) {};
+    } else {
+      SpillCalleeSavedRegisters = doSpillCalleeSavedRegs;
+      ScavengeFrameVirtualRegs = doScavengeFrameVirtualRegs;
+      UsesCalleeSaves = true;
+    }
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   MachineFunctionProperties getRequiredProperties() const override {
-    return MachineFunctionProperties().set(
-        MachineFunctionProperties::Property::AllVRegsAllocated);
+    MachineFunctionProperties MFP;
+    if (UsesCalleeSaves)
+      MFP.set(MachineFunctionProperties::Property::AllVRegsAllocated);
+    return MFP;
   }
 
   /// runOnMachineFunction - Insert prolog/epilog code and replace abstract
   /// frame indexes with appropriate references.
   ///
   bool runOnMachineFunction(MachineFunction &Fn) override;
 
 private:
+  std::function<void(MachineFunction &MF, RegScavenger *RS,
+                     unsigned &MinCSFrameIndex, unsigned &MaxCSFrameIndex,
+                     const MBBVector &SaveBlocks,
+                     const MBBVector &RestoreBlocks)>
+      SpillCalleeSavedRegisters;
+  std::function<void(MachineFunction &MF, RegScavenger *RS)>
+      ScavengeFrameVirtualRegs;
+
+  bool UsesCalleeSaves = false;
+
   RegScavenger *RS;
 
   // MinCSFrameIndex, MaxCSFrameIndex - Keeps the range of callee saved
   // stack frame indexes.
-  unsigned MinCSFrameIndex, MaxCSFrameIndex;
+  unsigned MinCSFrameIndex = std::numeric_limits<unsigned>::max();
+  unsigned MaxCSFrameIndex = 0;
 
   // Save and Restore blocks of the current function. Typically there is a
   // single save block, unless Windows EH funclets are involved.
-  SmallVector<MachineBasicBlock *, 1> SaveBlocks;
-  SmallVector<MachineBasicBlock *, 4> RestoreBlocks;
+  MBBVector SaveBlocks;
+  MBBVector RestoreBlocks;
 
   // Flag to control whether to use the register scavenger to resolve
   // frame index materialization registers. Set according to
   // TRI->requiresFrameIndexScavenging() for the current function.
   bool FrameIndexVirtualScavenging;
 
-  void calculateSets(MachineFunction &Fn);
-  void calculateCallsInformation(MachineFunction &Fn);
-  void assignCalleeSavedSpillSlots(MachineFunction &Fn,
-                                   const BitVector &SavedRegs);
-  void insertCSRSpillsAndRestores(MachineFunction &Fn);
+  void calculateCallFrameInfo(MachineFunction &Fn);
+  void calculateSaveRestoreBlocks(MachineFunction &Fn);
+
   void calculateFrameObjectOffsets(MachineFunction &Fn);
   void replaceFrameIndices(MachineFunction &Fn);
   void replaceFrameIndices(MachineBasicBlock *BB, MachineFunction &Fn,
                            int &SPAdj);
-  void scavengeFrameVirtualRegs(MachineFunction &Fn);
   void insertPrologEpilogCode(MachineFunction &Fn);
 };
 } // namespace
 
 char PEI::ID = 0;
 char &llvm::PrologEpilogCodeInserterID = PEI::ID;
 
 static cl::opt<unsigned>
 WarnStackSize("warn-stack-size", cl::Hidden, cl::init((unsigned)-1),
               cl::desc("Warn for stack size bigger than the given"
                        " number"));
 
-INITIALIZE_PASS_BEGIN(PEI, "prologepilog",
-                "Prologue/Epilogue Insertion", false, false)
+INITIALIZE_TM_PASS_BEGIN(PEI, "prologepilog", "Prologue/Epilogue Insertion",
+                         false, false)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
 INITIALIZE_PASS_DEPENDENCY(StackProtector)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
-INITIALIZE_PASS_END(PEI, "prologepilog",
+INITIALIZE_TM_PASS_END(PEI, "prologepilog",
                        "Prologue/Epilogue Insertion & Frame Finalization",
                        false, false)
 
+MachineFunctionPass *
+llvm::createPrologEpilogInserterPass(const TargetMachine *TM) {
+  return new PEI(TM);
+}
+
 STATISTIC(NumScavengedRegs, "Number of frame index regs scavenged");
 STATISTIC(NumBytesStackSpace,
           "Number of bytes used for stack in all functions");
 
 void PEI::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   AU.addPreserved<MachineLoopInfo>();
   AU.addPreserved<MachineDominatorTree>();
   AU.addRequired<StackProtector>();
   AU.addRequired<TargetPassConfig>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-/// Compute the set of return blocks
-void PEI::calculateSets(MachineFunction &Fn) {
-  const MachineFrameInfo *MFI = Fn.getFrameInfo();
-
-  // Even when we do not change any CSR, we still want to insert the
-  // prologue and epilogue of the function.
-  // So set the save points for those.
-
-  // Use the points found by shrink-wrapping, if any.
-  if (MFI->getSavePoint()) {
-    SaveBlocks.push_back(MFI->getSavePoint());
-    assert(MFI->getRestorePoint() && "Both restore and save must be set");
-    MachineBasicBlock *RestoreBlock = MFI->getRestorePoint();
-    // If RestoreBlock does not have any successor and is not a return block
-    // then the end point is unreachable and we do not need to insert any
-    // epilogue.
-    if (!RestoreBlock->succ_empty() || RestoreBlock->isReturnBlock())
-      RestoreBlocks.push_back(RestoreBlock);
-    return;
-  }
-
-  // Save refs to entry and return blocks.
-  SaveBlocks.push_back(&Fn.front());
-  for (MachineBasicBlock &MBB : Fn) {
-    if (MBB.isEHFuncletEntry())
-      SaveBlocks.push_back(&MBB);
-    if (MBB.isReturnBlock())
-      RestoreBlocks.push_back(&MBB);
-  }
-}
 
 /// StackObjSet - A set of stack object indexes
 typedef SmallSetVector<int, 8> StackObjSet;
 
 /// runOnMachineFunction - Insert prolog/epilog code and replace abstract
 /// frame indexes with appropriate references.
 ///
 bool PEI::runOnMachineFunction(MachineFunction &Fn) {
   const Function* F = Fn.getFunction();
   const TargetRegisterInfo *TRI = Fn.getSubtarget().getRegisterInfo();
   const TargetFrameLowering *TFI = Fn.getSubtarget().getFrameLowering();
 
-  assert(!Fn.getRegInfo().getNumVirtRegs() && "Regalloc must assign all vregs");
-
   RS = TRI->requiresRegisterScavenging(Fn) ? new RegScavenger() : nullptr;
   FrameIndexVirtualScavenging = TRI->requiresFrameIndexScavenging(Fn);
 
   // Calculate the MaxCallFrameSize and AdjustsStack variables for the
   // function's frame information. Also eliminates call frame pseudo
   // instructions.
-  calculateCallsInformation(Fn);
+  calculateCallFrameInfo(Fn);
 
-  // Determine which of the registers in the callee save list should be saved.
-  BitVector SavedRegs;
-  TFI->determineCalleeSaves(Fn, SavedRegs, RS);
-
-  // Insert spill code for any callee saved registers that are modified.
-  assignCalleeSavedSpillSlots(Fn, SavedRegs);
-
-  // Determine placement of CSR spill/restore code:
+  // Determine placement of CSR spill/restore code and prolog/epilog code:
   // place all spills in the entry block, all restores in return blocks.
-  calculateSets(Fn);
+  calculateSaveRestoreBlocks(Fn);
 
-  // Add the code to save and restore the callee saved registers.
-  if (!F->hasFnAttribute(Attribute::Naked))
-    insertCSRSpillsAndRestores(Fn);
+  // Handle CSR spilling and restoring, for targets that need it.
+  SpillCalleeSavedRegisters(Fn, RS, MinCSFrameIndex, MaxCSFrameIndex,
+                            SaveBlocks, RestoreBlocks);
 
   // Allow the target machine to make final modifications to the function
   // before the frame layout is finalized.
   TFI->processFunctionBeforeFrameFinalized(Fn, RS);
 
   // Calculate actual frame offsets for all abstract stack objects...
   calculateFrameObjectOffsets(Fn);
 
   // Add prolog and epilog code to the function.  This function is required
   // to align the stack frame as necessary for any stack variables or
   // called functions.  Because of this, calculateCalleeSavedRegisters()
   // must be called before this function in order to set the AdjustsStack
   // and MaxCallFrameSize variables.
   if (!F->hasFnAttribute(Attribute::Naked))
     insertPrologEpilogCode(Fn);
 
   // Replace all MO_FrameIndex operands with physical register references
   // and actual offsets.
   //
   replaceFrameIndices(Fn);
 
   // If register scavenging is needed, as we've enabled doing it as a
   // post-pass, scavenge the virtual registers that frame index elimination
   // inserted.
-  if (TRI->requiresRegisterScavenging(Fn) && FrameIndexVirtualScavenging)
-    scavengeFrameVirtualRegs(Fn);
+  if (TRI->requiresRegisterScavenging(Fn) && FrameIndexVirtualScavenging) {
+      ScavengeFrameVirtualRegs(Fn, RS);
 
       // Clear any vregs created by virtual scavenging.
       Fn.getRegInfo().clearVirtRegs();
+  }
 
   // Warn on stack size when we exceeds the given limit.
   MachineFrameInfo *MFI = Fn.getFrameInfo();
   uint64_t StackSize = MFI->getStackSize();
   if (WarnStackSize.getNumOccurrences() > 0 && WarnStackSize < StackSize) {
     DiagnosticInfoStackSize DiagStackSize(*F, StackSize);
     F->getContext().diagnose(DiagStackSize);
   }
 
   delete RS;
   SaveBlocks.clear();
   RestoreBlocks.clear();
   return true;
 }
 
-/// calculateCallsInformation - Calculate the MaxCallFrameSize and AdjustsStack
+/// Calculate the MaxCallFrameSize and AdjustsStack
 /// variables for the function's frame information and eliminate call frame
 /// pseudo instructions.
-void PEI::calculateCallsInformation(MachineFunction &Fn) {
+void PEI::calculateCallFrameInfo(MachineFunction &Fn) {
   const TargetInstrInfo &TII = *Fn.getSubtarget().getInstrInfo();
   const TargetFrameLowering *TFI = Fn.getSubtarget().getFrameLowering();
   MachineFrameInfo *MFI = Fn.getFrameInfo();
 
   unsigned MaxCallFrameSize = 0;
   bool AdjustsStack = MFI->adjustsStack();
 
   // Get the function call frame set-up and tear-down instruction opcode
for (std::vector<MachineBasicBlock::iterator>::iterator
     // the target doesn't indicate otherwise, remove the call frame pseudos
     // here. The sub/add sp instruction pairs are still inserted, but we don't
     // need to track the SP adjustment for frame index elimination.
     if (TFI->canSimplifyCallFramePseudos(Fn))
       TFI->eliminateCallFramePseudoInstr(Fn, *I->getParent(), I);
   }
 }
 
-void PEI::assignCalleeSavedSpillSlots(MachineFunction &F,
-                                      const BitVector &SavedRegs) {
-  // These are used to keep track the callee-save area. Initialize them.
-  MinCSFrameIndex = INT_MAX;
-  MaxCSFrameIndex = 0;
+/// Compute the sets of entry and return blocks for saving and restoring
+/// callee-saved registers, and placing prolog and epilog code.
+void PEI::calculateSaveRestoreBlocks(MachineFunction &Fn) {
+  const MachineFrameInfo *MFI = Fn.getFrameInfo();
+
+  // Even when we do not change any CSR, we still want to insert the
+  // prologue and epilogue of the function.
+  // So set the save points for those.
 
+  // Use the points found by shrink-wrapping, if any.
+  if (MFI->getSavePoint()) {
+    SaveBlocks.push_back(MFI->getSavePoint());
+    assert(MFI->getRestorePoint() && "Both restore and save must be set");
+    MachineBasicBlock *RestoreBlock = MFI->getRestorePoint();
+    // If RestoreBlock does not have any successor and is not a return block
+    // then the end point is unreachable and we do not need to insert any
+    // epilogue.
+    if (!RestoreBlock->succ_empty() || RestoreBlock->isReturnBlock())
+      RestoreBlocks.push_back(RestoreBlock);
+    return;
+  }
+
+  // Save refs to entry and return blocks.
+  SaveBlocks.push_back(&Fn.front());
+  for (MachineBasicBlock &MBB : Fn) {
+    if (MBB.isEHFuncletEntry())
+      SaveBlocks.push_back(&MBB);
+    if (MBB.isReturnBlock())
+      RestoreBlocks.push_back(&MBB);
+  }
+}
+
+static void assignCalleeSavedSpillSlots(MachineFunction &F,
+                                        const BitVector &SavedRegs,
+                                        unsigned &MinCSFrameIndex,
+                                        unsigned &MaxCSFrameIndex) {
   if (SavedRegs.empty())
     return;
 
   const TargetRegisterInfo *RegInfo = F.getSubtarget().getRegisterInfo();
   const MCPhysReg *CSRegs = RegInfo->getCalleeSavedRegs(&F);
 
   std::vector<CalleeSavedInfo> CSI;
   for (unsigned i = 0; CSRegs[i]; ++i) {
for (MachineBasicBlock *MBB : Visited) {
         MBB->addLiveIn(Reg);
     }
   }
 }
 
 /// insertCSRSpillsAndRestores - Insert spill and restore code for
 /// callee saved registers used in the function.
 ///
-void PEI::insertCSRSpillsAndRestores(MachineFunction &Fn) {
+static void insertCSRSpillsAndRestores(MachineFunction &Fn,
+                                       const MBBVector &SaveBlocks,
+                                       const MBBVector &RestoreBlocks) {
   // Get callee saved register information.
   MachineFrameInfo *MFI = Fn.getFrameInfo();
   const std::vector<CalleeSavedInfo> &CSI = MFI->getCalleeSavedInfo();
 
   MFI->setCalleeSavedInfoValid(true);
 
   // Early exit if no callee saved registers are modified!
   if (CSI.empty())
if (!TFI->restoreCalleeSavedRegisters(*MBB, I, CSI, TRI)) {
           I = BeforeI;
           ++I;
         }
       }
     }
   }
 }
 
+static void doSpillCalleeSavedRegs(MachineFunction &Fn, RegScavenger *RS,
+                                   unsigned &MinCSFrameIndex,
+                                   unsigned &MaxCSFrameIndex,
+                                   const MBBVector &SaveBlocks,
+                                   const MBBVector &RestoreBlocks) {
+  const Function *F = Fn.getFunction();
+  const TargetFrameLowering *TFI = Fn.getSubtarget().getFrameLowering();
+  MinCSFrameIndex = std::numeric_limits<unsigned>::max();
+  MaxCSFrameIndex = 0;
+
+  // Determine which of the registers in the callee save list should be saved.
+  BitVector SavedRegs;
+  TFI->determineCalleeSaves(Fn, SavedRegs, RS);
+
+  // Assign stack slots for any callee-saved registers that must be spilled.
+  assignCalleeSavedSpillSlots(Fn, SavedRegs, MinCSFrameIndex, MaxCSFrameIndex);
+
+  // Add the code to save and restore the callee saved registers.
+  if (!F->hasFnAttribute(Attribute::Naked))
+    insertCSRSpillsAndRestores(Fn, SaveBlocks, RestoreBlocks);
+}
+
 /// AdjustStackOffset - Helper function used to adjust the stack frame offset.
 static inline void
 AdjustStackOffset(MachineFrameInfo *MFI, int FrameIdx,
                   bool StackGrowsDown, int64_t &Offset,
                   unsigned &MaxAlign, unsigned Skew) {
   // If the stack grows down, add the object size to find the lowest address.
   if (StackGrowsDown)
     Offset += MFI->getObjectSize(FrameIdx);
for (unsigned i = MinCSFrameIndex; i <= MaxCSFrameIndex; ++i) {
       unsigned Align = MFI->getObjectAlignment(i);
       // Adjust to alignment boundary
       Offset = alignTo(Offset, Align, Skew);
 
       DEBUG(dbgs() << "alloc FI(" << i << ") at SP[" << -Offset << "]\n");
       MFI->setObjectOffset(i, -Offset);        // Set the computed offset
     }
   } else {
-    int MaxCSFI = MaxCSFrameIndex, MinCSFI = MinCSFrameIndex;
-    for (int i = MaxCSFI; i >= MinCSFI ; --i) {
+    unsigned MaxCSFI = MaxCSFrameIndex, MinCSFI = MinCSFrameIndex;
+    for (unsigned i = MaxCSFI; i >= MinCSFI; --i) {
       unsigned Align = MFI->getObjectAlignment(i);
       // Adjust to alignment boundary
       Offset = alignTo(Offset, Align, Skew);
 
       DEBUG(dbgs() << "alloc FI(" << i << ") at SP[" << Offset << "]\n");
       MFI->setObjectOffset(i, Offset);
       Offset += MFI->getObjectSize(i);
     }
for (MachineBasicBlock::iterator I = BB->begin(); I != BB->end(); ) {
 
     if (DoIncr && I != BB->end()) ++I;
 
     // Update register states.
     if (RS && !FrameIndexVirtualScavenging && MI) RS->forward(MI);
   }
 }
 
-/// scavengeFrameVirtualRegs - Replace all frame index virtual registers
+/// doScavengeFrameVirtualRegs - Replace all frame index virtual registers
 /// with physical registers. Use the register scavenger to find an
 /// appropriate register to use.
 ///
 /// FIXME: Iterating over the instruction stream is unnecessary. We can simply
 /// iterate over the vreg use list, which at this point only contains machine
 /// operands for which eliminateFrameIndex need a new scratch reg.
-void
-PEI::scavengeFrameVirtualRegs(MachineFunction &Fn) {
+static void
+doScavengeFrameVirtualRegs(MachineFunction &Fn, RegScavenger *RS) {
   // Run through the instructions and find any virtual registers.
   for (MachineFunction::iterator BB = Fn.begin(),
        E = Fn.end(); BB != E; ++BB) {
     RS->enterBasicBlock(*BB);
 
     int SPAdj = 0;
 
     // The instruction stream may change in the loop, so check BB->end()

lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

 
 void WebAssemblyPassConfig::addPostRegAlloc() {
   // TODO: The following CodeGen passes don't currently support code containing
   // virtual registers. Consider removing their restrictions and re-enabling
   // them.
 
   // Has no asserts of its own, but was not written to handle virtual regs.
   disablePass(&ShrinkWrapID);
-  // We use our own PrologEpilogInserter which is very slightly modified to
-  // tolerate virtual registers.
-  disablePass(&PrologEpilogCodeInserterID);
 
   // These functions all require the AllVRegsAllocated property.
   disablePass(&MachineCopyPropagationID);
   disablePass(&PostRASchedulerID);
   disablePass(&FuncletLayoutID);
   disablePass(&StackMapLivenessID);
   disablePass(&LiveDebugValuesID);
 
   if (getOptLevel() != CodeGenOpt::None) {
     // Mark registers as representing wasm's expression stack.
     addPass(createWebAssemblyRegStackify());
 
     // Run the register coloring pass to reduce the total number of registers.
     addPass(createWebAssemblyRegColoring());
   }
 
   TargetPassConfig::addPostRegAlloc();
-
-  // Run WebAssembly's version of the PrologEpilogInserter. Target-independent
-  // PEI runs after PostRegAlloc and after ShrinkWrap. Putting it here will run
-  // PEI before ShrinkWrap but otherwise in the same position in the order.
-  addPass(createWebAssemblyPEI());
 }
 
 void WebAssemblyPassConfig::addPreEmitPass() {
   TargetPassConfig::addPreEmitPass();
 
   // Eliminate multiple-entry loops.
   addPass(createWebAssemblyFixIrreducibleControlFlow());