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

148

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

148

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.

60

Indeed :).

200

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

625

Unrelated change.

[dschuff]

lib/CodeGen/PrologEpilogInserter.cpp

200

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

625

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

60

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

60

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.

63

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

63

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

63

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!

lib/CodeGen/PrologEpilogInserter.cpp

 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include <climits>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "pei"
 
 namespace {
 class PEI : public MachineFunctionPass {
 public:
   static char ID;
   PEI() : MachineFunctionPass(ID) {
     initializePEIPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   /// runOnMachineFunction - Insert prolog/epilog code and replace abstract
   /// frame indexes with appropriate references.
   ///
   bool runOnMachineFunction(MachineFunction &Fn) override;
 
 private:
   RegScavenger *RS;
 
   // MinCSFrameIndex, MaxCSFrameIndex - Keeps the range of callee saved
   // stack frame indexes.
   unsigned MinCSFrameIndex, MaxCSFrameIndex;
 
   // 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;
 
   // 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 calculateCallFrameInfo(MachineFunction &Fn);
+  void calculateSaveRestoreBlocks(MachineFunction &Fn);
+
   void assignCalleeSavedSpillSlots(MachineFunction &Fn,
                                    const BitVector &SavedRegs);
   void insertCSRSpillsAndRestores(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
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);
 
+  // Virtual-register machines such as WebAssembly do not use normal register
+  // allocation. Such targets also do not use callee-save registers and do
+  // not require register scavenging.
+  const bool HasVirtualRegisters = Fn.getRegInfo().getNumVirtRegs();
+
   // Calculate the MaxCallFrameSize and AdjustsStack variables for the
   // function's frame information. Also eliminates call frame pseudo
   // instructions.
-  calculateCallsInformation(Fn);
+  calculateCallFrameInfo(Fn);
+
+  // Determine placement of CSR spill/restore code and prolog/epilog code:
+  // place all spills in the entry block, all restores in return blocks.
+  calculateSaveRestoreBlocks(Fn);
 
+  // Handle CSR spilling and restoring, for targets that need it.
+  if (HasVirtualRegisters) {
+    assert((!TRI->getCalleeSavedRegs(&Fn) ||
+            TRI->getCalleeSavedRegs(&Fn)[0] == 0) &&
+           "All registers should be allocated for non-virtual targets");
+  } else {
     // 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.
+    // Assign stack slots for any callee-saved registers that must be spilled.
     assignCalleeSavedSpillSlots(Fn, SavedRegs);
 
-  // Determine placement of CSR spill/restore code:
-  // place all spills in the entry block, all restores in return blocks.
-  calculateSets(Fn);
-
     // Add the code to save and restore the callee saved registers.
     if (!F->hasFnAttribute(Attribute::Naked))
       insertCSRSpillsAndRestores(Fn);
+  }
 
   // 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 (HasVirtualRegisters) {
+    assert(!RS && !TRI->requiresFrameIndexScavenging(Fn) &&
+           "Virtual-register targets do not support register scavenging");
+  } else {
     // 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);
 
     // 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);
   }
 }
 
+/// 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);
+  }
+}
+
+
 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;
 
   if (SavedRegs.empty())
     return;
for (unsigned i = MinCSFrameIndex; i <= MaxCSFrameIndex; ++i) {
       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 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);
     }
   }