Merge rust-lang#48

48: Add function prologue info to stackmaps. r=vext01 a=ptersilie



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

Author: bors[bot]

llvm/include/llvm/CodeGen/StackMaps.h

@@ -38,7 +38,7 @@ class StackMapOpers {
   enum { IDPos, NBytesPos };
 
 private:
-  const MachineInstr* MI;
+  const MachineInstr *MI;
 
 public:
   explicit StackMapOpers(const MachineInstr *MI);
@@ -204,6 +204,12 @@ class StackMaps {
         : Reg(Reg), DwarfRegNum(DwarfRegNum), Size(Size) {}
   };
 
+  // Stores callee-saved register info.
+  struct CSR {
+    int Reg = 0;
+    int Offset = 0;
+  };
+
   // OpTypes are used to encode information about the following logical
   // operand (which may consist of several MachineOperands) for the
   // OpParser.
@@ -237,13 +243,19 @@ class StackMaps {
   using LiveVarsVec = SmallVector<LocationVec, 8>;
   using LiveOutVec = SmallVector<LiveOutReg, 8>;
   using ConstantPool = MapVector<uint64_t, uint64_t>;
+  using CSRVec = SmallVector<CSR, 8>;
 
   struct FunctionInfo {
     uint64_t StackSize = 0;
     uint64_t RecordCount = 1;
+    bool HasFramePointer;
+    CSRVec SpilledRegisters;
 
     FunctionInfo() = default;
-    explicit FunctionInfo(uint64_t StackSize) : StackSize(StackSize) {}
+    explicit FunctionInfo(uint64_t StackSize, bool HasFramePointer,
+                          CSRVec SpilledRegisters)
+        : StackSize(StackSize), HasFramePointer(HasFramePointer),
+          SpilledRegisters(SpilledRegisters) {}
   };
 
   struct CallsiteInfo {
@@ -265,16 +277,13 @@ 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);
 
   /// Generate a stackmap record for a patchpoint instruction.
-  void recordPatchPoint(const MCSymbol &L,
-                        const MachineInstr &MI);
+  void recordPatchPoint(const MCSymbol &L, const MachineInstr &MI);
 
   /// Generate a stackmap record for a statepoint instruction.
-  void recordStatepoint(const MCSymbol &L,
-                        const MachineInstr &MI);
+  void recordStatepoint(const MCSymbol &L, const MachineInstr &MI);
 
   /// If there is any stack map data, create a stack map section and serialize
   /// the map info into it. This clears the stack map data structures
@@ -294,6 +303,7 @@ class StackMaps {
   CallsiteInfoList CSInfos;
   ConstantPool ConstPool;
   FnInfoMap FnInfos;
+  bool HasFramePointer;
 
   MachineInstr::const_mop_iterator
   parseOperand(MachineInstr::const_mop_iterator MOI,
@@ -321,10 +331,9 @@ class StackMaps {
   /// STACKMAP, and PATCHPOINT the label is expected to immediately *preceed*
   /// lowering of the MI to MCInsts.  For STATEPOINT, it expected to
   /// immediately *follow*.  It's not clear this difference was intentional,
-  /// but it exists today.  
-  void recordStackMapOpers(const MCSymbol &L,
-                           const MachineInstr &MI, uint64_t ID,
-                           MachineInstr::const_mop_iterator MOI,
+  /// but it exists today.
+  void recordStackMapOpers(const MCSymbol &L, const MachineInstr &MI,
+                           uint64_t ID, MachineInstr::const_mop_iterator MOI,
                            MachineInstr::const_mop_iterator MOE,
                            bool recordResult = false);
 
@@ -340,6 +349,9 @@ class StackMaps {
   /// Emit the callsite info for each stackmap/patchpoint intrinsic call.
   void emitCallsiteEntries(MCStreamer &OS);
 
+  /// Emit information about the function prologue.
+  void emitCSRInfo(MCStreamer &OS);
+
   void print(raw_ostream &OS);
   void debug() { print(dbgs()); }
 };

llvm/include/llvm/Transforms/Yk/ControlPoint.h

@@ -10,6 +10,10 @@
 // The name of the new control point replacing the user's dummy control point.
 #define YK_NEW_CONTROL_POINT "__ykrt_control_point"
 
+// The name of the function which reconstructs the stackframe and jumps to the
+// right instruction in AOT from where to continue.
+#define YK_RECONSTRUCT_FRAMES "__ykrt_reconstruct_frames"
+
 namespace llvm {
 ModulePass *createYkControlPointPass();
 } // namespace llvm

llvm/lib/CodeGen/StackMaps.cpp

@@ -15,6 +15,7 @@
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
@@ -54,10 +55,8 @@ static uint64_t getConstMetaVal(const MachineInstr &MI, unsigned Idx) {
   return MO.getImm();
 }
 
-StackMapOpers::StackMapOpers(const MachineInstr *MI)
-  : MI(MI) {
-  assert(getVarIdx() <= MI->getNumOperands() &&
-         "invalid stackmap definition");
+StackMapOpers::StackMapOpers(const MachineInstr *MI) : MI(MI) {
+  assert(getVarIdx() <= MI->getNumOperands() && "invalid stackmap definition");
 }
 
 PatchPointOpers::PatchPointOpers(const MachineInstr *MI)
@@ -81,11 +80,10 @@ unsigned PatchPointOpers::getNextScratchIdx(unsigned StartIdx) const {
 
   // Find the next scratch register (implicit def and early clobber)
   unsigned ScratchIdx = StartIdx, e = MI->getNumOperands();
-  while (ScratchIdx < e &&
-         !(MI->getOperand(ScratchIdx).isReg() &&
-           MI->getOperand(ScratchIdx).isDef() &&
-           MI->getOperand(ScratchIdx).isImplicit() &&
-           MI->getOperand(ScratchIdx).isEarlyClobber()))
+  while (ScratchIdx < e && !(MI->getOperand(ScratchIdx).isReg() &&
+                             MI->getOperand(ScratchIdx).isDef() &&
+                             MI->getOperand(ScratchIdx).isImplicit() &&
+                             MI->getOperand(ScratchIdx).isEarlyClobber()))
     ++ScratchIdx;
 
   assert(ScratchIdx != e && "No scratch register available");
@@ -274,8 +272,8 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
     if (SubRegIdx)
       Offset = TRI->getSubRegIdxOffset(SubRegIdx);
 
-    Locs.emplace_back(Location::Register, TRI->getSpillSize(*RC),
-                      DwarfRegNum, Offset);
+    Locs.emplace_back(Location::Register, TRI->getSpillSize(*RC), DwarfRegNum,
+                      Offset);
     return ++MOI;
   }
 
@@ -570,6 +568,7 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
 
   // Record the stack size of the current function and update callsite count.
   const MachineFrameInfo &MFI = AP.MF->getFrameInfo();
+  const TargetRegisterInfo *TRI = AP.MF->getSubtarget().getRegisterInfo();
   const TargetRegisterInfo *RegInfo = AP.MF->getSubtarget().getRegisterInfo();
   bool HasDynamicFrameSize =
       MFI.hasVarSizedObjects() || RegInfo->hasStackRealignment(*(AP.MF));
@@ -578,17 +577,30 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
   auto CurrentIt = FnInfos.find(AP.CurrentFnSym);
   if (CurrentIt != FnInfos.end())
     CurrentIt->second.RecordCount++;
-  else
-    FnInfos.insert(std::make_pair(AP.CurrentFnSym, FunctionInfo(FrameSize)));
+  else {
+    // Collect callee-saved-register spills.
+    CSRVec CSRInfo;
+    std::vector<CalleeSavedInfo> CSI = MFI.getCalleeSavedInfo();
+    for (auto cs : CSI) {
+      int dwreg = getDwarfRegNum(cs.getReg(), TRI);
+      CSR Entry = {dwreg, cs.getFrameIdx()};
+      CSRInfo.push_back(Entry);
+    }
+    // Check whether this function pushed a frame pointer.
+    const TargetFrameLowering *TFL = AP.MF->getSubtarget().getFrameLowering();
+    bool HasFramePointer = TFL->hasFP(*(AP.MF));
+    FnInfos.insert(std::make_pair(
+        AP.CurrentFnSym, FunctionInfo(FrameSize, HasFramePointer, CSRInfo)));
+  }
 }
 
 void StackMaps::recordStackMap(const MCSymbol &L, const MachineInstr &MI) {
   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()),
+  recordStackMapOpers(L, MI, ID,
+                      std::next(MI.operands_begin(), opers.getVarIdx()),
                       MI.operands_end());
 }
 
@@ -784,6 +796,32 @@ void StackMaps::emitCallsiteEntries(MCStreamer &OS) {
   }
 }
 
+// Emit information about the function prologue: pushed frame pointer and
+// callee-saved registers.
+// CSRInfo[NumFunctions] {
+//     uint8  : HasFramePtr
+//     uint8  : Padding
+//     uint32 : NumSpills
+//
+//     Spills[NumSpills] {
+//         uint16 : DwarfRegister
+//     	   uint16 : Padding
+//     	   int32  : Offset
+//     }
+// }
+void StackMaps::emitCSRInfo(MCStreamer &OS) {
+  for (auto const &FR : FnInfos) {
+    OS.emitInt8(FR.second.HasFramePointer);
+    OS.emitInt8(0); // Padding
+    OS.emitInt32(FR.second.SpilledRegisters.size());
+    for (auto &Ent : FR.second.SpilledRegisters) {
+      OS.emitInt16(Ent.Reg);
+      OS.emitInt16(0); // Padding
+      OS.emitInt32(Ent.Offset);
+    }
+  }
+}
+
 /// Serialize the stackmap data.
 void StackMaps::serializeToStackMapSection() {
   (void)WSMP;
@@ -812,6 +850,7 @@ void StackMaps::serializeToStackMapSection() {
   emitFunctionFrameRecords(OS);
   emitConstantPoolEntries(OS);
   emitCallsiteEntries(OS);
+  emitCSRInfo(OS);
   OS.addBlankLine();
 
   // Clean up.

llvm/lib/Transforms/Yk/ControlPoint.cpp

@@ -162,25 +162,11 @@ class YkControlPoint : public ModulePass {
             ->getType();
 
     // Create the new control point, which is of the form:
-    //   bool new_control_point(YkMT*, YkLocation*, CtrlPointVars*,
-    //   ReturnValue*)
-    // If the return type of the control point's caller is void (i.e. if a
-    // function f calls yk_control_point and f's return type is void), create
-    // an Int1 pointer as a dummy. We have to pass something as the yk_stopgap
-    // signature expects a pointer, even if its never used.
-    Type *ReturnTy = Caller->getReturnType();
-    Type *ReturnPtrTy;
-    if (ReturnTy->isVoidTy()) {
-      // Create dummy pointer which we pass in but which is never written to.
-      ReturnPtrTy = Type::getInt1Ty(Context);
-    } else {
-      ReturnPtrTy = ReturnTy;
-    }
-    FunctionType *FType =
-        FunctionType::get(Type::getInt1Ty(Context),
-                          {YkMTTy, YkLocTy, CtrlPointVarsTy->getPointerTo(),
-                           ReturnPtrTy->getPointerTo()},
-                          false);
+    //   void* new_control_point(YkMT*, YkLocation*, CtrlPointVars*, void*)
+    PointerType *VoidPtr = PointerType::get(Context, 0);
+    FunctionType *FType = FunctionType::get(
+        VoidPtr, {YkMTTy, YkLocTy, CtrlPointVarsTy->getPointerTo(), VoidPtr},
+        false);
     Function *NF = Function::Create(FType, GlobalVariable::ExternalLinkage,
                                     YK_NEW_CONTROL_POINT, M);
 
@@ -190,10 +176,6 @@ class YkControlPoint : public ModulePass {
     IRBuilder<> Builder(Caller->getEntryBlock().getFirstNonPHI());
     Value *InputStruct = Builder.CreateAlloca(CtrlPointVarsTy, 0, "");
 
-    // Also at the top, generate storage for the interpreted return of the
-    // control points caller.
-    Value *ReturnPtr = Builder.CreateAlloca(ReturnPtrTy, 0, "");
-
     Builder.SetInsertPoint(OldCtrlPointCall);
     unsigned LvIdx = 0;
     for (Value *LV : LiveVals) {
@@ -205,11 +187,17 @@ class YkControlPoint : public ModulePass {
       LvIdx++;
     }
 
+    // Create a call to the llvm.frameaddress intrinsic, which we pass into the
+    // control point. This is later required for stack reconstruction.
+    Function *FrameAddress =
+        Intrinsic::getDeclaration(&M, Intrinsic::frameaddress, {VoidPtr});
+    Value *FAddr = Builder.CreateCall(FrameAddress, {Builder.getInt32(0)});
+
     // Insert call to the new control point.
     Instruction *NewCtrlPointCallInst = Builder.CreateCall(
         NF, {OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_MT_IDX),
              OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_LOC_IDX),
-             InputStruct, ReturnPtr});
+             InputStruct, FAddr});
 
     // Once the control point returns we need to extract the (potentially
     // mutated) values from the returned YkCtrlPointStruct and reassign them to
@@ -238,14 +226,14 @@ class YkControlPoint : public ModulePass {
     // Create the new exit block.
     BasicBlock *ExitBB = BasicBlock::Create(Context, "", Caller);
     Builder.SetInsertPoint(ExitBB);
-    // YKFIXME: We need to return the value of interpreted return and the return
-    // type must be that of the control point's caller.
-    if (ReturnTy->isVoidTy()) {
-      Builder.CreateRetVoid();
-    } else {
-      Value *ReturnValue = Builder.CreateLoad(ReturnTy, ReturnPtr);
-      Builder.CreateRet(ReturnValue);
-    }
+
+    // Create call to frame reconstructor.
+    FunctionType *YKFRType =
+        FunctionType::get(Type::getVoidTy(Context), {VoidPtr}, false);
+    Function *YKFR = Function::Create(YKFRType, GlobalVariable::ExternalLinkage,
+                                      YK_RECONSTRUCT_FRAMES, M);
+    Builder.CreateCall(YKFR, {NewCtrlPointCallInst});
+    Builder.CreateUnreachable();
 
     // To do so we need to first split up the current block and then
     // insert a conditional branch that either continues or returns.
@@ -256,7 +244,13 @@ class YkControlPoint : public ModulePass {
     Instruction &OldBr = BB->back();
     OldBr.eraseFromParent();
     Builder.SetInsertPoint(BB);
-    Builder.CreateCondBr(NewCtrlPointCallInst, ExitBB, ContBB);
+    // The return value of the control point tells us whether a guard has failed
+    // or not (i.e. anything other than a nullptr means a guard has failed). If
+    // it has jump to `ExitBB` which calls the code that copies over the new
+    // stack from the given pointer.
+    Value *HasGuardFailed = Builder.CreateICmpEQ(
+        NewCtrlPointCallInst, ConstantPointerNull::get(VoidPtr));
+    Builder.CreateCondBr(HasGuardFailed, ContBB, ExitBB);
 
 #ifndef NDEBUG
     // Our pass runs after LLVM normally does its verify pass. In debug builds

llvm/test/CodeGen/X86/statepoint-stackmap-size.ll

@@ -4,7 +4,7 @@
 ; spaces) starting with a `.` follow (e.g. `  .byte`).
 ;
 ;      CHECK:	.section	.llvm_stackmaps,{{.*$}}
-; CHECK-NEXT:{{(.+$[[:space:]]){51}[[:space:]]}}
+; CHECK-NEXT:{{(.+$[[:space:]]){54}[[:space:]]}}
 ;  CHECK-NOT:{{.|[[:space:]]}}
 
 target triple = "x86_64-pc-linux-gnu"