Merge rust-lang#57

57: Enable stackmap support for tracking multiple locations. r=ltratt a=ptersilie



Co-authored-by: Lukas Diekmann <[email protected]>

Author: bors[bot]

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);
 }

llvm/lib/Transforms/Yk/StackMaps.cpp

@@ -56,9 +56,10 @@ class YkStackmaps : public ModulePass {
             CallInst &CI = cast<CallInst>(I);
             if (CI.isInlineAsm())
               continue;
-            if (CI.isIndirectCall())
-              continue;
-            if (CI.getCalledFunction()->isIntrinsic())
+            // We don't need to insert stackmaps after intrinsics. But since we
+            // can't tell if an indirect call is an intrinsic at compile time,
+            // emit a stackmap in those cases too.
+            if (!CI.isIndirectCall() && CI.getCalledFunction()->isIntrinsic())
               continue;
             SMCalls.insert({&I, LA.getLiveVarsBefore(&I)});
           } else if ((isa<BranchInst>(I) &&