Merge rust-lang#52

52: Fix stackmaps. r=ltratt a=ptersilie



Co-authored-by: Lukas Diekmann <lukas.diekmann@gmail.com>

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -329,6 +329,12 @@ class AsmPrinter : public MachineFunctionPass {
   /// definition in the same module.
   MCSymbol *getSymbolPreferLocal(const GlobalValue &GV) const;
 
+  /// Maps call instructions to the label emitted just after the call. The label
+  /// is later used to calculate the correct offset of the call instruction
+  /// without interference from caller-saved registers which are sometimes
+  /// emitted inbetween call and stackmap.
+  MCSymbol *YkLastCallLabel = nullptr;
+
   //===------------------------------------------------------------------===//
   // XRay instrumentation implementation.
   //===------------------------------------------------------------------===//

llvm/include/llvm/Transforms/Yk/LivenessAnalysis.h

@@ -1,6 +1,7 @@
 #ifndef __YK_LIVENESS_H
 #define __YK_LIVENESS_H
 
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 
 #include <map>
@@ -22,15 +23,18 @@ class LivenessAnalysis {
   /// Find the successor instructions of the specified instruction.
   std::set<Instruction *> getSuccessorInstructions(Instruction *I);
 
+  // A domniator tree used to sort the live variables.
+  DominatorTree DT;
+
   /// Replaces the set `S` with the set `R`, returning if the set changed.
   bool updateValueSet(std::set<Value *> &S, const std::set<Value *> &R);
 
 public:
   LivenessAnalysis(Function *Func);
 
-  /// Returns the set of live variables immediately before the specified
-  /// instruction.
-  std::set<Value *> getLiveVarsBefore(Instruction *I);
+  /// Returns the vector of live variables immediately before the specified
+  /// instruction (sorted in order of appearance).
+  std::vector<Value *> getLiveVarsBefore(Instruction *I);
 };
 
 } // namespace llvm

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -135,6 +135,7 @@ static cl::opt<bool>
 
 extern bool YkAllocLLVMBBAddrMapSection;
 extern bool YkExtendedLLVMBBAddrMapSection;
+extern bool YkStackMapOffsetFix;
 
 const char DWARFGroupName[] = "dwarf";
 const char DWARFGroupDescription[] = "DWARF Emission";
@@ -1699,6 +1700,9 @@ void AsmPrinter::emitFunctionBody() {
   bool HasAnyRealCode = false;
   int NumInstsInFunction = 0;
 
+  // Reset YkLastCallLabel so we don't use it across different functions.
+  YkLastCallLabel = nullptr;
+
   bool CanDoExtraAnalysis = ORE->allowExtraAnalysis(DEBUG_TYPE);
   for (auto &MBB : *MF) {
     // Print a label for the basic block.
@@ -1800,6 +1804,31 @@ void AsmPrinter::emitFunctionBody() {
         }
 
         emitInstruction(&MI);
+        // Generate labels for function calls so we can record the correct
+        // instruction offset. The conditions for generating the label must be
+        // the same as the ones for generating the stackmap call in
+        // `Transforms/Yk/Stackmaps.cpp`, as otherwise we could end up with
+        // wrong offsets (e.g. we create a label here but the corresponding
+        // stackmap call was ommitted, and this label is then used for the
+        // following stackmap call).
+
+        // Convoluted way of finding our whether this function call is an
+        // intrinsic.
+        bool IsIntrinsic = false;
+        if (MI.getNumExplicitDefs() < MI.getNumOperands()) {
+          IsIntrinsic = MI.getOperand(MI.getNumExplicitDefs()).isIntrinsicID();
+        }
+        if (YkStackMapOffsetFix && MI.isCall() && !MI.isInlineAsm() &&
+            (MI.getOpcode() != TargetOpcode::STACKMAP) &&
+            (MI.getOpcode() != TargetOpcode::PATCHPOINT) && !IsIntrinsic) {
+          // YKFIXME: We don't need to emit labels (and stackmap calls) for
+          // functions that cannot guard fail, e.g. inlined functions, or
+          // functions we don't have IR for.
+          MCSymbol *MILabel =
+              OutStreamer->getContext().createTempSymbol("ykstackcall", true);
+          OutStreamer->emitLabel(MILabel);
+          YkLastCallLabel = MILabel;
+        }
         if (CanDoExtraAnalysis) {
           MCInst MCI;
           MCI.setOpcode(MI.getOpcode());

llvm/lib/Support/Yk.cpp

@@ -18,3 +18,10 @@ static cl::opt<bool, true> YkExtendedLLVMBBAddrMapSectionParser(
     "yk-extended-llvmbbaddrmap-section",
     cl::desc("Use the extended Yk `.llvmbbaddrmap` section format"),
     cl::NotHidden, cl::location(YkExtendedLLVMBBAddrMapSection));
+
+bool YkStackMapOffsetFix;
+static cl::opt<bool, true> YkStackMapOffsetFixParser(
+    "yk-stackmap-offset-fix",
+    cl::desc("Apply a fix to stackmaps that corrects the reported instruction "
+             "offset in the presence of calls."),
+    cl::NotHidden, cl::location(YkStackMapOffsetFix));

llvm/lib/Target/X86/X86MCInstLower.cpp

@@ -1422,10 +1422,25 @@ void X86AsmPrinter::LowerPATCHABLE_OP(const MachineInstr &MI,
 void X86AsmPrinter::LowerSTACKMAP(const MachineInstr &MI) {
   SMShadowTracker.emitShadowPadding(*OutStreamer, getSubtargetInfo());
 
-  auto &Ctx = OutStreamer->getContext();
-  MCSymbol *MILabel = Ctx.createTempSymbol();
-  OutStreamer->emitLabel(MILabel);
-
+  // When the stackmap is attached to a function call, caller-saved register
+  // loads may be emitted in-between call and stackmap call, thus skewing the
+  // reported offset.
+  // This makes it impossible to reliably continue at the correct location after
+  // deoptimisation has occured. In order to fix this, we insert labels after
+  // every call that could lead to a guard failure, and store the last such
+  // label inside `YkLastCallLabel`. If `YkLastCallLabel` is set (i.e. the
+  // stackmap was inserted after a call) use it to calculate the stackmap
+  // offset, otherwise (i.e. the stackmap was inserted before a branch) generate
+  // a new label as per ususual.
+  MCSymbol *MILabel;
+  if (YkLastCallLabel != nullptr) {
+    MILabel = YkLastCallLabel;
+    YkLastCallLabel = nullptr;
+  } else {
+    auto &Ctx = OutStreamer->getContext();
+    MILabel = Ctx.createTempSymbol();
+    OutStreamer->emitLabel(MILabel);
+  }
   SM.recordStackMap(*MILabel, MI);
   unsigned NumShadowBytes = MI.getOperand(1).getImm();
   SMShadowTracker.reset(NumShadowBytes);

llvm/lib/Transforms/Yk/ControlPoint.cpp

@@ -130,7 +130,8 @@ class YkControlPoint : public ModulePass {
 
     // Find all live variables just before the call to the control point.
     LivenessAnalysis LA(OldCtrlPointCall->getFunction());
-    const std::set<Value *> LiveVals = LA.getLiveVarsBefore(OldCtrlPointCall);
+    const std::vector<Value *> LiveVals =
+        LA.getLiveVarsBefore(OldCtrlPointCall);
     if (LiveVals.size() == 0) {
       Context.emitError(
           "The interpreter loop has no live variables!\n"

llvm/lib/Transforms/Yk/LivenessAnalysis.cpp

@@ -69,6 +69,10 @@ LivenessAnalysis::LivenessAnalysis(Function *Func) {
   // Compute defs and uses for each instruction.
   std::map<Instruction *, std::set<Value *>> Defs;
   std::map<Instruction *, std::set<Value *>> Uses;
+
+  // Create a domniator tree so we can later sort the live variables.
+  DT = DominatorTree(*Func);
+
   for (BasicBlock &BB : *Func) {
     for (Instruction &I : BB) {
       // Record what this instruction defines.
@@ -172,8 +176,19 @@ LivenessAnalysis::LivenessAnalysis(Function *Func) {
   } while (Changed); // Until a fixed-point.
 }
 
-std::set<Value *> LivenessAnalysis::getLiveVarsBefore(Instruction *I) {
-  return In[I];
+std::vector<Value *> LivenessAnalysis::getLiveVarsBefore(Instruction *I) {
+  std::set<Value *> Res = In[I];
+  // Sort the live variables by order of appearance using a dominator tree. The
+  // order is important for frame construction during deoptimisation: since live
+  // variables may reference other live variables they need to be proceesed in
+  // the order they appear in the module.
+  std::vector<Value *> Sorted(Res.begin(), Res.end());
+  std::sort(Sorted.begin(), Sorted.end(), [this](Value *A, Value *B) {
+    if (isa<Instruction>(B))
+      return DT.dominates(A, cast<Instruction>(B));
+    return false;
+  });
+  return Sorted;
 }
 
 } // namespace llvm

llvm/lib/Transforms/Yk/StackMaps.cpp

@@ -45,17 +45,28 @@ class YkStackmaps : public ModulePass {
       return false;
     }
 
-    std::map<Instruction *, std::set<Value *>> SMCalls;
+    std::map<Instruction *, std::vector<Value *>> SMCalls;
     for (Function &F : M) {
       if (F.empty()) // skip declarations.
         continue;
       LivenessAnalysis LA(&F);
       for (BasicBlock &BB : F)
-        for (Instruction &I : BB)
-          if ((isa<CallInst>(I)) ||
-              ((isa<BranchInst>(I)) && (cast<BranchInst>(I).isConditional())) ||
-              isa<SwitchInst>(I))
+        for (Instruction &I : BB) {
+          if (isa<CallInst>(I)) {
+            CallInst &CI = cast<CallInst>(I);
+            if (CI.isInlineAsm())
+              continue;
+            if (CI.isIndirectCall())
+              continue;
+            if (CI.getCalledFunction()->isIntrinsic())
+              continue;
             SMCalls.insert({&I, LA.getLiveVarsBefore(&I)});
+          } else if ((isa<BranchInst>(I) &&
+                      cast<BranchInst>(I).isConditional()) ||
+                     isa<SwitchInst>(I)) {
+            SMCalls.insert({&I, LA.getLiveVarsBefore(&I)});
+          }
+        }
     }
 
     Function *SMFunc = Intrinsic::getDeclaration(&M, SMFuncID);
@@ -65,14 +76,20 @@ class YkStackmaps : public ModulePass {
     Value *Shadow = ConstantInt::get(Type::getInt32Ty(Context), 0);
     for (auto It : SMCalls) {
       Instruction *I = cast<Instruction>(It.first);
-      const std::set<Value *> L = It.second;
+      const std::vector<Value *> L = It.second;
 
       IRBuilder<> Bldr(I);
       Value *SMID = ConstantInt::get(Type::getInt64Ty(Context), Count);
       std::vector<Value *> Args = {SMID, Shadow};
       for (Value *A : L)
         Args.push_back(A);
 
+      if (isa<CallInst>(I)) {
+        // Insert the stackmap call after (not before) the call instruction, so
+        // the offset of the stackmap entry will record the instruction after
+        // the call, which is where we want to continue after deoptimisation.
+        Bldr.SetInsertPoint(I->getNextNonDebugInstruction());
+      }
       Bldr.CreateCall(SMFunc->getFunctionType(), SMFunc,
                       ArrayRef<Value *>(Args));
       Count++;