Emit label after calls to fix stackmap offsets.

During codegen stackmaps instructions emit a label when they are
encountered. This label is then used to calculate the offset of its
location from function entry. Unfortunately, this offset isn't always
correct since codegen may insert instructions (e.g. for restoring
caller-saved registers) in-between a call and the stackmap instruction.
To fix this, we now insert a label immediately after codegenning a call.
The following stackmap instruction can then pickup that label to compute
the correct offset.

Author: ptersilie

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/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/StackMaps.cpp

@@ -51,11 +51,22 @@ class YkStackmaps : public ModulePass {
         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);