Track more than one location per variable in stackmaps.

Stackmaps only track one location per live variable, e.g. in a register,
on the stack, etc. However, when deoptimising to native stack frames
values may exist in multiple locations, e.g. in a register as well as on
the stack or in two registers.

The main reason for this, at least from observation, are spills before
function calls. For example, if a value is held in $rdi and that
register needs to be freed up to execute a call, the register allocator
may spill it to a register or onto the stack. After the call is done the
value is reloaded. A simple optimisation is to cache the spills so the
value can be reloaded throughout a function without having to respill it
every time. Depending on where the stackmap call is located we can only
see the restored register but not the spill. During deoptimisation we
then only populate the tracked register with the right value. When the
program then decides to restore the register again (e.g. because it has
executed another call) loading it from the stack won't recover the value
since we haven't pushed it during optimisation.

This commit adds an analysis to the compiler chain that attempts to find
additional locations for live variables and encode them in stackmaps. It
does so but traversing the control flow of a function and creating a
mapping on any register moves, stack stores or reloads. For simplicity
(and so we don't need to change the stackmap format) we assume that only
one such additional location may exist (which may very well turn out to
be wrong in the future).

These changes enable us to run much further in some of the lua tests,
though there are still some bugs to be fixed.

Author: ptersilie

llvm/include/llvm/CodeGen/StackMaps.h

@@ -285,7 +285,7 @@ class StackMaps {
   /// Generate a stackmap record for a stackmap instruction.
   ///
   /// MI must be a raw STACKMAP, not a PATCHPOINT.
-  void recordStackMap(const MCSymbol &L, const MachineInstr &MI);
+  void recordStackMap(const MCSymbol &L, const MachineInstr &MI, std::map<Register, int64_t> SpillsOffsets = {});
 
   /// Generate a stackmap record for a patchpoint instruction.
   void recordPatchPoint(const MCSymbol &L, const MachineInstr &MI);
@@ -316,7 +316,8 @@ class StackMaps {
   MachineInstr::const_mop_iterator
   parseOperand(MachineInstr::const_mop_iterator MOI,
                MachineInstr::const_mop_iterator MOE, LiveVarsVec &LiveVars,
-               LiveOutVec &LiveOuts) const;
+               LiveOutVec &LiveOuts,
+               std::map<Register, int64_t> SpillOffsets = {}) const;
 
   /// Specialized parser of statepoint operands.
   /// They do not directly correspond to StackMap record entries.
@@ -343,6 +344,7 @@ class StackMaps {
   void recordStackMapOpers(const MCSymbol &L, const MachineInstr &MI,
                            uint64_t ID, MachineInstr::const_mop_iterator MOI,
                            MachineInstr::const_mop_iterator MOE,
+                           std::map<Register, int64_t> SpillOffsets = {},
                            bool recordResult = false);
 
   /// Emit the stackmap header.

llvm/lib/CodeGen/StackMaps.cpp

@@ -18,6 +18,7 @@
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/MC/MCContext.h"
@@ -219,7 +220,8 @@ static unsigned getDwarfRegNum(unsigned Reg, const TargetRegisterInfo *TRI) {
 MachineInstr::const_mop_iterator
 StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
                         MachineInstr::const_mop_iterator MOE,
-                        LiveVarsVec &LiveVars, LiveOutVec &LiveOuts) const {
+                        LiveVarsVec &LiveVars, LiveOutVec &LiveOuts,
+                        std::map<Register, int64_t> SpillOffsets) const {
   LocationVec &Locs = LiveVars.back();
   const TargetRegisterInfo *TRI = AP.MF->getSubtarget().getRegisterInfo();
   if (MOI->isImm()) {
@@ -282,7 +284,27 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
     const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(MOI->getReg());
     assert(!MOI->getSubReg() && "Physical subreg still around.");
 
-    unsigned Offset = 0;
+
+    signed Offset = 0;
+    // Check if there are any mappings for this register in the spillmap. If so,
+    // encode the additional location in the offset field of the stackmap record
+    // (which is unused for register locations). Note that this assumes that
+    // there can only be one additional location for each value, which may turn
+    // out to be false.
+    if (MOI->isReg()) {
+      Register R = MOI->getReg();
+      if (SpillOffsets.count(R) > 0) {
+        Offset = SpillOffsets[R];
+        assert(SpillOffsets[R] != 0);
+        if (Offset > 0) {
+          // If the additional location is another register encode its DWARF id.
+          // Also temporarily add 1 since 0 is used to mean there is no
+          // additional location.
+          Offset = getDwarfRegNum(Offset, TRI) + 1;
+        }
+      }
+    }
+
     unsigned DwarfRegNum = getDwarfRegNum(MOI->getReg(), TRI);
     unsigned LLVMRegNum = *TRI->getLLVMRegNum(DwarfRegNum, false);
     unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, MOI->getReg());
@@ -520,6 +542,7 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
                                     const MachineInstr &MI, uint64_t ID,
                                     MachineInstr::const_mop_iterator MOI,
                                     MachineInstr::const_mop_iterator MOE,
+                                    std::map<Register, int64_t> SpillOffsets,
                                     bool recordResult) {
   MCContext &OutContext = AP.OutStreamer->getContext();
 
@@ -529,7 +552,7 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
   if (recordResult) {
     assert(PatchPointOpers(&MI).hasDef() && "Stackmap has no return value.");
     parseOperand(MI.operands_begin(), std::next(MI.operands_begin()), LiveVars,
-                 LiveOuts);
+                 LiveOuts, SpillOffsets);
     LiveVars.push_back(LocationVec());
   }
 
@@ -538,7 +561,7 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
     parseStatepointOpers(MI, MOI, MOE, LiveVars, LiveOuts);
   else
     while (MOI != MOE)
-      MOI = parseOperand(MOI, MOE, LiveVars, LiveOuts);
+      MOI = parseOperand(MOI, MOE, LiveVars, LiveOuts, SpillOffsets);
 
   // Move large constants into the constant pool.
   for (auto &Locations : LiveVars) {
@@ -611,14 +634,14 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
   }
 }
 
-void StackMaps::recordStackMap(const MCSymbol &L, const MachineInstr &MI) {
+void StackMaps::recordStackMap(const MCSymbol &L, const MachineInstr &MI, std::map<Register, int64_t> SpillOffsets) {
   assert(MI.getOpcode() == TargetOpcode::STACKMAP && "expected stackmap");
 
   StackMapOpers opers(&MI);
   const int64_t ID = MI.getOperand(PatchPointOpers::IDPos).getImm();
   recordStackMapOpers(L, MI, ID,
                       std::next(MI.operands_begin(), opers.getVarIdx()),
-                      MI.operands_end());
+                      MI.operands_end(), SpillOffsets);
 }
 
 void StackMaps::recordPatchPoint(const MCSymbol &L, const MachineInstr &MI) {
@@ -627,7 +650,7 @@ void StackMaps::recordPatchPoint(const MCSymbol &L, const MachineInstr &MI) {
   PatchPointOpers opers(&MI);
   const int64_t ID = opers.getID();
   auto MOI = std::next(MI.operands_begin(), opers.getStackMapStartIdx());
-  recordStackMapOpers(L, MI, ID, MOI, MI.operands_end(),
+  recordStackMapOpers(L, MI, ID, MOI, MI.operands_end(), {},
                       opers.isAnyReg() && opers.hasDef());
 
 #ifndef NDEBUG
@@ -648,7 +671,7 @@ void StackMaps::recordStatepoint(const MCSymbol &L, const MachineInstr &MI) {
   StatepointOpers opers(&MI);
   const unsigned StartIdx = opers.getVarIdx();
   recordStackMapOpers(L, MI, opers.getID(), MI.operands_begin() + StartIdx,
-                      MI.operands_end(), false);
+                      MI.operands_end(), {}, false);
 }
 
 /// Emit the stackmap header.

llvm/lib/Support/Yk.cpp

@@ -25,3 +25,9 @@ static cl::opt<bool, true> YkStackMapOffsetFixParser(
     cl::desc("Apply a fix to stackmaps that corrects the reported instruction "
              "offset in the presence of calls."),
     cl::NotHidden, cl::location(YkStackMapOffsetFix));
+
+bool YkStackMapAdditionalLocs;
+static cl::opt<bool, true> YkStackMapAdditionalLocsParser(
+    "yk-stackmap-add-locs",
+    cl::desc("Encode additional locations for registers into stackmaps."),
+    cl::NotHidden, cl::location(YkStackMapAdditionalLocs));

llvm/lib/Target/X86/X86AsmPrinter.cpp

@@ -46,6 +46,7 @@
 #include "llvm/Target/TargetMachine.h"
 
 using namespace llvm;
+extern bool YkStackMapAdditionalLocs;
 
 X86AsmPrinter::X86AsmPrinter(TargetMachine &TM,
                              std::unique_ptr<MCStreamer> Streamer)
@@ -55,6 +56,116 @@ X86AsmPrinter::X86AsmPrinter(TargetMachine &TM,
 // Primitive Helper Functions.
 //===----------------------------------------------------------------------===//
 
+const TargetInstrInfo *TII;
+
+/// Clear any mappings that map to the given register.
+void clearRhs(Register Reg, std::map<Register, int64_t> &SpillMap) {
+  auto I = SpillMap.begin();
+  while (I != SpillMap.end()) {
+    if (I->second == Reg) {
+      I = SpillMap.erase(I);
+    } else {
+      ++I;
+    }
+  }
+}
+
+/// Given a MachineBasicBlock, analyse its instructions to build a mapping of
+/// where duplicate values live. This can be either in another register or on
+/// the stack. Since registers are always positive and stack offsets negative,
+/// we can store both in the same map while still being able to distinguish
+/// them. This allows us to keep the changes in Stackmaps.cpp to a minimum and
+/// avoids having to alter the stackmap format.
+void processInstructions(
+    const MachineBasicBlock *MBB,
+    std::map<Register, int64_t> &SpillMap,
+    std::map<const MachineInstr *, std::map<Register, int64_t>> &StackmapSpillMaps
+  ) {
+  for (const MachineInstr &Instr : MBB->instrs()) {
+    // At each stackmap call, save the current mapping so it can later be
+    // encoded in the stackmap when it is lowered.
+    if (Instr.getOpcode() == TargetOpcode::STACKMAP) {
+      StackmapSpillMaps[&Instr] = SpillMap;
+      continue;
+    }
+
+    // Copying a value from one register B to another A, creates a mapping from
+    // A to B. If A is tracked by the stackmap, then B will also be tracked and
+    // assigned the same value during deoptimisation.
+    if (Instr.isMoveReg()) {
+      const MachineOperand Lhs = Instr.getOperand(0);
+      const MachineOperand Rhs = Instr.getOperand(1);
+      assert(Lhs.isReg() && "Is register.");
+      assert(Rhs.isReg() && "Is register.");
+      SpillMap[Lhs.getReg()] = Rhs.getReg();
+      // Reassigning a new value to Lhs means any mappings to Lhs are now void
+      // and need to be removed.
+      clearRhs(Lhs.getReg(), SpillMap);
+      continue;
+    }
+
+    int FI;
+    // A value from a register was spilled to the stack. Create a mapping from
+    // that register to the stack offset.
+    if (TII->isStoreToStackSlotPostFE(Instr, FI)) {
+      const MachineOperand OffsetOp = Instr.getOperand(3);
+      const MachineOperand MO = Instr.getOperand(5);
+      assert(MO.isReg() && "Is register.");
+      const Register Reg = MO.getReg();
+      if (OffsetOp.isImm()) {
+        const int64_t Offset = OffsetOp.getImm();
+        SpillMap[Reg] = Offset;
+        clearRhs(Reg, SpillMap);
+      }
+      continue;
+    }
+
+    // A value from the stack was loaded back into a register. Create a mapping from
+    // that register to the stack offset.
+    if (TII->isLoadFromStackSlotPostFE(Instr, FI)) {
+      const MachineOperand OffsetOp = Instr.getOperand(4);
+      const MachineOperand Lhs = Instr.getOperand(0);
+      assert(Lhs.isReg() && "Is register.");
+      const Register Reg = Lhs.getReg();
+      if (OffsetOp.isImm()) {
+        const int64_t Offset = OffsetOp.getImm();
+        SpillMap[Reg] = Offset;
+        clearRhs(Reg, SpillMap);
+      }
+      continue;
+    }
+
+    // Any other assignments to tracked registers removes their mapping.
+    for (const MachineOperand MO : Instr.defs()) {
+      assert(MO.isReg() && "Is register.");
+      SpillMap.erase(MO.getReg());
+      clearRhs(MO.getReg(), SpillMap);
+    }
+  }
+}
+
+/// Walk through the control flow of a function starting at the entry block
+/// using depth first search. This makes sure the mappings are are generated in
+/// the correct order.
+/// YKFIXME: Can be updated to use an iterative approach.
+void findSpillLocations(
+    const MachineBasicBlock *MBB,
+    std::set<const MachineBasicBlock *> &Seen,
+    std::map<Register, int64_t> SpillMap,
+    std::map<const MachineInstr *, std::map<Register, int64_t>> &StackmapSpillMaps
+    ) {
+
+  // Block has already been processed.
+  if (Seen.count(MBB) > 0) {
+    return;
+  }
+  Seen.insert(MBB);
+  processInstructions(MBB, SpillMap, StackmapSpillMaps);
+  for (MachineBasicBlock *Succ : MBB->successors()) {
+    findSpillLocations(Succ, Seen, SpillMap, StackmapSpillMaps);
+  }
+}
+
 /// runOnMachineFunction - Emit the function body.
 ///
 bool X86AsmPrinter::runOnMachineFunction(MachineFunction &MF) {
@@ -82,6 +193,15 @@ bool X86AsmPrinter::runOnMachineFunction(MachineFunction &MF) {
     OutStreamer->endCOFFSymbolDef();
   }
 
+  // Analyse control flow to find out the location of register spills to the
+  // stack or to other registers, so we can also encode those in stackmaps.
+  TII = MF.getSubtarget().getInstrInfo();
+  std::set<const MachineBasicBlock *> S;
+  std::map<Register, int64_t> SpillMap;
+  if (YkStackMapAdditionalLocs) {
+    findSpillLocations(&*MF.begin(), S, SpillMap, StackmapSpillMaps);
+  }
+
   // Emit the rest of the function body.
   emitFunctionBody();
 

llvm/lib/Target/X86/X86AsmPrinter.h

@@ -71,6 +71,8 @@ class LLVM_LIBRARY_VISIBILITY X86AsmPrinter : public AsmPrinter {
   };
 
   StackMapShadowTracker SMShadowTracker;
+  // Map holding information about register spills.
+  std::map<const MachineInstr *, std::map<Register, int64_t>> StackmapSpillMaps;
 
   // All instructions emitted by the X86AsmPrinter should use this helper
   // method.

llvm/lib/Target/X86/X86MCInstLower.cpp

@@ -1502,7 +1502,7 @@ void X86AsmPrinter::LowerSTACKMAP(const MachineInstr &MI) {
     MILabel = Ctx.createTempSymbol();
     OutStreamer->emitLabel(MILabel);
   }
-  SM.recordStackMap(*MILabel, MI);
+  SM.recordStackMap(*MILabel, MI, StackmapSpillMaps[&MI]);
   unsigned NumShadowBytes = MI.getOperand(1).getImm();
   SMShadowTracker.reset(NumShadowBytes);
 }