Add support for symbolic large constant entries inside stackmaps

[undingen]

Summary

This patch implements support for embedding symbols in the array of large constants inside a stackmap/patchpoint.
We (=Pyston project) embed in the LLVM IR we generate a lot of pointers to external constants. I would like to cache the JITed functions and therefore I need to make the embedded constants relocatable -> change the IR to refer to module variables instead which I can replace when loading the cached object.
Currently when a patchpoint has a constant variable as live argument it will materialize it in a register. Which we don't want because the live variables are rarely used (also the regalloc runs out of registers when having to many const variables as argument to the pp?!?... I filled PR23306).
This patch emits a symbol reference inside the array of large constants in the stackmap. The stack map format is untouched by this patch.

Together with http://lists.cs.uiuc.edu/pipermail/llvm-commits/Week-of-Mon-20150330/269160.html this should make it possible to cache more JITed functions containing patchpoints.

[reames]

Comments inline, but you are still missing the most important part of the patch. Where is the documentation which describes the output format? I have inferred some of it from the code, but it will save you time if we can settle on the output before hashing through every bit of code. We may need to change the code based on the result of the format discussion.

One question to consider: As a consumer of the generated stackmap, how do I tell a given entry is a actual constant vs a symbol that needs resolved? Maybe I'm missing the obvious, but I didn't see this in the code. Or are you assuming that the consumer is only parsing the finalized/full relocated version of the binary section?

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

6920

I don't understand this check. What case is this for? Should this possibly be an assert?

Or should you possibly be using the getValueType in place of the MVT::i64 as an argument to getTargetGlobalAddress?

lib/CodeGen/StackMaps.cpp

309

The need for this check makes me thing that you should introduce a new Location type.

329

I really think you want a Location::GlobalValue or something similar.

333

This seems unnecessarily complicated. Just emit the offset from the beginning of the global value section. (This comment is dependent on the one below w.r.t. reuse of constant section.)

420

Mixing global values with constants doesn't seem like the right approach here. Why not introduce new section specifically for global values?

(I'm open to being convinced this is the right answer, but you need to make the argument.)

[undingen]

Thanks for the comments and sorry for not having written any documentation because I thought the default behavior of "RTDyldMemoryManager::getSymbolAddress" which will resolve the symbols makes documenting the change unnecessary but I can see that this may be unexpected and needs documentation.
I reused the existing stackmap format in order to not break existing clients. That's why I choose to put them in the constant value table but this means it's impossible for clients to retrieve the number of large constants which don't need symbol resolving and which do. (They will just be at the end of the table but there is no way to know how many there are at the end)
If changing the stackmap format is not a big problem than maybe it's best to make a new section, which would make the resulting stackmap and the code to generate it inside llvm more obvious.

The MVT::i64 check is my attempt at making sure that the symbol addresses are exactly 64bit long like the entries in the stackmap constant table. But this may should get changed to allow smaller values.

Thanks again for the comments, I will wait with changing the patch until we all agreed on how to proceed.

In D9176#161717, @reames wrote:

Comments inline, but you are still missing the most important part of the patch. Where is the documentation which describes the output format? I have inferred some of it from the code, but it will save you time if we can settle on the output before hashing through every bit of code. We may need to change the code based on the result of the format discussion.

One question to consider: As a consumer of the generated stackmap, how do I tell a given entry is a actual constant vs a symbol that needs resolved? Maybe I'm missing the obvious, but I didn't see this in the code. Or are you assuming that the consumer is only parsing the finalized/full relocated version of the binary section?

[reames]

In D9176#161738, @undingen wrote:

I reused the existing stackmap format in order to not break existing clients. That's why I choose to put them in the constant value table but this means it's impossible for clients to retrieve the number of large constants which don't need symbol resolving and which do. (They will just be at the end of the table but there is no way to know how many there are at the end)

I would strongly prefer you extended the format. Reusing the existing sections for mixed purposes without updating the documentation would be a really bad idea.

If changing the stackmap format is not a big problem than maybe it's best to make a new section, which would make the resulting stackmap and the code to generate it inside llvm more obvious.

The section format is versioned. Adding a new section is not a big deal.

The MVT::i64 check is my attempt at making sure that the symbol addresses are exactly 64bit long like the entries in the stackmap constant table. But this may should get changed to allow smaller values.

This is an incompatible assumption with non-x86 bit architectures right? I'd prefer not to see this baked in.

Thanks again for the comments, I will wait with changing the patch until we all agreed on how to proceed.

I would suggest that you update *just the docs* for the next couple of iterations. Until we stablize on what the feature will look like, getting the code exactly right will be a waste of time.

I say this specifically because I don't feel like I understand exactly what you're trying to propose. Until that part is clear, I can't really offer anything in the way of useful review.

[undingen]

In D9176#188933, @reames wrote:

I say this specifically because I don't feel like I understand exactly what you're trying to propose. Until that part is clear, I can't really offer anything in the way of useful review.

Sorry I overlooked your reply.
Maybe I can make clear what I have in mind by describing how I'm currently using this patch:

We (Pyston project) use patchpoints to implement inline caches and for deoptimization when using the LLVM tier.
For the deoptimization use case we add all variables to the patchpoint live args which we need too continue the execution in a lower generic tier (e.g. interpreter). A lot of our generated IR values were direct inttoptr casts because we often generate instances of our objects outside of LLVM. For example we may generate instances of a python objects when we setup the internal representation of a python function which we then share between the interpreter and LLVM tier. That's why we had a lot of inttoptr casts in our generated IR, there are also additional args like pointers to the AST nodes which we will need for deopt.

Deopimizations should happen only very rarely that means that we don't want to actually load all the constants we specified as live values inside the patchpoint into registers/stack slots. Currently LLVM will put all arguments which are constants inside the stackmap constant table in order to not have to generate code in front of the patchpoint to put all this constant values into register/stack slots. This is exactly how I would expect the behavior to be and how I need it.

But then I added a new feature: in order to speedup JITing time if we encounter the same function on the next application start I implemented an object cache for the LLVM generated code. This means I need to be able to relocate all this embedded pointers because the memory layout will not be the same. I choose to solve this by emitting special unique symbol name for all cases where I previously embedded the direct pointer value. This symbol names are deterministic, on the next startup when encountering the same function I can directly load it from the object cache and just have to return the real pointer values inside the RTDyldMemoryManager::getSymbolAddress() overloaded function.

The problem I encountered and this patch tries to solve is that LLVM will currently emit code which will load all this symbolic constants into registers before the patchpoint. With this patch we will stop emitting this machine instructions and instead emit constant table entries inside the stackmap.

Hope this helps understanding what I have done (even if my english isn't good), I successfully use this solution now since several weeks and it gave us a huge speedup.

include/llvm/CodeGen/StackMaps.h

 public:
   struct Location {
     enum LocationType { Unprocessed, Register, Direct, Indirect, Constant,
                         ConstantIndex };
     LocationType LocType;
     unsigned Size;
     unsigned Reg;
     int64_t Offset;
-    Location() : LocType(Unprocessed), Size(0), Reg(0), Offset(0) {}
+    const MCSymbol *Sym;
+    Location() : LocType(Unprocessed), Size(0), Reg(0), Offset(0), Sym(0) {}
     Location(LocationType LocType, unsigned Size, unsigned Reg, int64_t Offset)
-      : LocType(LocType), Size(Size), Reg(Reg), Offset(Offset) {}
+        : LocType(LocType), Size(Size), Reg(Reg), Offset(Offset), Sym(0) {}
+    Location(const MCSymbol *Sym)
+        : LocType(LocationType::Constant), Size(sizeof(int64_t)), Reg(0),
+          Offset(0), Sym(Sym) {}
   };
 
   struct LiveOutReg {
     unsigned short Reg;
     unsigned short RegNo;
     unsigned short Size;
 
     LiveOutReg() : Reg(0), RegNo(0), Size(0) {}
     LiveOutReg(unsigned short Reg, unsigned short RegNo, unsigned short Size)
       : Reg(Reg), RegNo(RegNo), Size(Size) {}
 
     void MarkInvalid() { Reg = 0; }
 
     // Only sort by the dwarf register number.
     bool operator< (const LiveOutReg &LO) const { return RegNo < LO.RegNo; }
     static bool IsInvalid(const LiveOutReg &LO) { return LO.Reg == 0; }
   };
 
   // OpTypes are used to encode information about the following logical
   // operand (which may consist of several MachineOperands) for the
   // OpParser.
-  typedef enum { DirectMemRefOp, IndirectMemRefOp, ConstantOp } OpType;
+  typedef enum {
+    DirectMemRefOp,
+    IndirectMemRefOp,
+    ConstantOp,
+    ConstantGVOp
+  } OpType;
 
   StackMaps(AsmPrinter &AP);
 
   void reset() {
     CSInfos.clear();
     ConstPool.clear();
+    ConstSymPool.clear();
     FnStackSize.clear();
   }
 
   /// \brief Generate a stackmap record for a stackmap instruction.
   ///
   /// MI must be a raw STACKMAP, not a PATCHPOINT.
   void recordStackMap(const MachineInstr &MI);
 
   /// \brief Generate a stackmap record for a patchpoint instruction.
   void recordPatchPoint(const MachineInstr &MI);
 
   /// \brief Generate a stackmap record for a statepoint instruction.
   void recordStatepoint(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
   /// afterwards.
   void serializeToStackMapSection();
 
 private:
   static const char *WSMP;
   typedef SmallVector<Location, 8> LocationVec;
   typedef SmallVector<LiveOutReg, 8> LiveOutVec;
   typedef MapVector<uint64_t, uint64_t> ConstantPool;
+  typedef MapVector<const MCSymbol *, const MCSymbol *> ConstantSymMap;
   typedef MapVector<const MCSymbol *, uint64_t> FnStackSizeMap;
 
   struct CallsiteInfo {
     const MCExpr *CSOffsetExpr;
     uint64_t ID;
     LocationVec Locations;
     LiveOutVec LiveOuts;
     CallsiteInfo() : CSOffsetExpr(nullptr), ID(0) {}
     CallsiteInfo(const MCExpr *CSOffsetExpr, uint64_t ID,
                  LocationVec &&Locations, LiveOutVec &&LiveOuts)
       : CSOffsetExpr(CSOffsetExpr), ID(ID), Locations(std::move(Locations)),
         LiveOuts(std::move(LiveOuts)) {}
   };
 
   typedef std::vector<CallsiteInfo> CallsiteInfoList;
 
   AsmPrinter &AP;
   CallsiteInfoList CSInfos;
   ConstantPool ConstPool;
+  ConstantSymMap ConstSymPool;
   FnStackSizeMap FnStackSize;
 
   MachineInstr::const_mop_iterator
   parseOperand(MachineInstr::const_mop_iterator MOI,
                MachineInstr::const_mop_iterator MOE,
                LocationVec &Locs, LiveOutVec &LiveOuts) const;
 
   /// \brief Create a live-out register record for the given register @p Reg.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(OpVal)) {
       Ops.push_back(
         Builder.DAG.getTargetConstant(StackMaps::ConstantOp, MVT::i64));
       Ops.push_back(
         Builder.DAG.getTargetConstant(C->getSExtValue(), MVT::i64));
     } else if (FrameIndexSDNode *FI = dyn_cast<FrameIndexSDNode>(OpVal)) {
       const TargetLowering &TLI = Builder.DAG.getTargetLoweringInfo();
       Ops.push_back(
         Builder.DAG.getTargetFrameIndex(FI->getIndex(), TLI.getPointerTy()));
+    } else if (auto *GA = dyn_cast<GlobalAddressSDNode>(OpVal)) {
+      if (GA->getValueType(0) != MVT::i64)
+        Ops.push_back(OpVal);
+      else {
+        Ops.push_back(
+            Builder.DAG.getTargetConstant(StackMaps::ConstantGVOp, MVT::i64));
+        Ops.push_back(Builder.DAG.getTargetGlobalAddress(GA->getGlobal(),
+                                                         SDLoc(GA), MVT::i64));
+      }
     } else
       Ops.push_back(OpVal);
   }
 }
 
 /// \brief Lower llvm.experimental.stackmap directly to its target opcode.
 void SelectionDAGBuilder::visitStackmap(const CallInst &CI) {
   // void @llvm.experimental.stackmap(i32 <id>, i32 <numShadowBytes>,

lib/CodeGen/StackMaps.cpp

if (MOI->isImm()) {
     }
     case StackMaps::ConstantOp: {
       ++MOI;
       assert(MOI->isImm() && "Expected constant operand.");
       int64_t Imm = MOI->getImm();
       Locs.push_back(Location(Location::Constant, sizeof(int64_t), 0, Imm));
       break;
     }
+    case StackMaps::ConstantGVOp: {
+      ++MOI;
+      assert(MOI->isGlobal() && "Expected a global value operand.");
+      const GlobalValue *GV = MOI->getGlobal();
+      assert(GV);
+      MCSymbol *Sym = AP.TM.getSymbol(GV, *AP.Mang);
+      assert(Sym);
+      Locs.push_back(Location(Sym));
+      break;
+    }
     }
     return ++MOI;
   }
 
   // The physical register number will ultimately be encoded as a DWARF regno.
   // The stack map also records the size of a spill slot that can hold the
   // register content. (The runtime can track the actual size of the data type
   // if it needs to.)
while (MOI != MOE) {
     MOI = parseOperand(MOI, MOE, Locations, LiveOuts);
   }
 
   // Move large constants into the constant pool.
   for (LocationVec::iterator I = Locations.begin(), E = Locations.end();
        I != E; ++I) {
     // Constants are encoded as sign-extended integers.
     // -1 is directly encoded as .long 0xFFFFFFFF with no constant pool.
-    if (I->LocType == Location::Constant && !isInt<32>(I->Offset)) {
+    if (I->LocType == Location::Constant && !isInt<32>(I->Offset) && !I->Sym) {
       I->LocType = Location::ConstantIndex;
       // ConstPool is intentionally a MapVector of 'uint64_t's (as
       // opposed to 'int64_t's).  We should never be in a situation
       // where we have to insert either the tombstone or the empty
       // keys into a map, and for a DenseMap<uint64_t, T> these are
       // (uint64_t)0 and (uint64_t)-1.  They can be and are
       // represented using 32 bit integers.
 
       assert((uint64_t)I->Offset != DenseMapInfo<uint64_t>::getEmptyKey() &&
              (uint64_t)I->Offset != DenseMapInfo<uint64_t>::getTombstoneKey() &&
              "empty and tombstone keys should fit in 32 bits!");
       auto Result = ConstPool.insert(std::make_pair(I->Offset, I->Offset));
       I->Offset = Result.first - ConstPool.begin();
     }
   }
 
+  // Convert constant symbols to ConstantIndex entries.
+  for (LocationVec::iterator I = Locations.begin(), E = Locations.end(); I != E;
+       ++I) {
+    if (I->LocType == Location::Constant && I->Sym) {
+      I->LocType = Location::ConstantIndex;
+      auto Result = ConstSymPool.insert(std::make_pair(I->Sym, I->Sym));
+      // The symbolic entries will be emitted after the ConstPool entries.
+      I->Offset = ConstPool.size() + Result.first - ConstSymPool.begin();
+    }
+  }
+
   // Create an expression to calculate the offset of the callsite from function
   // entry.
   const MCExpr *CSOffsetExpr = MCBinaryExpr::CreateSub(
     MCSymbolRefExpr::Create(MILabel, OutContext),
     MCSymbolRefExpr::Create(AP.CurrentFnSymForSize, OutContext),
     OutContext);
 
   CSInfos.emplace_back(CSOffsetExpr, ID, std::move(Locations),
void StackMaps::emitStackmapHeader(MCStreamer &OS) {
   OS.EmitIntValue(StackMapVersion, 1); // Version.
   OS.EmitIntValue(0, 1); // Reserved.
   OS.EmitIntValue(0, 2); // Reserved.
 
   // Num functions.
   DEBUG(dbgs() << WSMP << "#functions = " << FnStackSize.size() << '\n');
   OS.EmitIntValue(FnStackSize.size(), 4);
   // Num constants.
-  DEBUG(dbgs() << WSMP << "#constants = " << ConstPool.size() << '\n');
-  OS.EmitIntValue(ConstPool.size(), 4);
+  auto NumConst = ConstPool.size() + ConstSymPool.size();
+  DEBUG(dbgs() << WSMP << "#constants = " << NumConst << '\n');
+  OS.EmitIntValue(NumConst, 4);
   // Num callsites.
   DEBUG(dbgs() << WSMP << "#callsites = " << CSInfos.size() << '\n');
   OS.EmitIntValue(CSInfos.size(), 4);
 }
 
 /// Emit the function frame record for each function.
 ///
 /// StkSizeRecord[NumFunctions] {

 /// int64  : Constants[NumConstants]
 void StackMaps::emitConstantPoolEntries(MCStreamer &OS) {
   // Constant pool entries.
   DEBUG(dbgs() << WSMP << "constants:\n");
   for (auto ConstEntry : ConstPool) {
     DEBUG(dbgs() << WSMP << ConstEntry.second << '\n');
     OS.EmitIntValue(ConstEntry.second, 8);
   }
+  for (auto ConstEntry : ConstSymPool) {
+    DEBUG(dbgs() << WSMP << ConstEntry.second << '\n');
+    OS.EmitSymbolValue(ConstEntry.second, 8);
+  }
 }
 
 /// Emit the callsite info for each callsite.
 ///
 /// StkMapRecord[NumRecords] {
 ///   uint64 : PatchPoint ID
 ///   uint32 : Instruction Offset
 ///   uint16 : Reserved (record flags)
for (const auto &CSI : CSInfos) {
     OS.EmitValueToAlignment(8);
   }
 }
 
 /// Serialize the stackmap data.
 void StackMaps::serializeToStackMapSection() {
   (void) WSMP;
   // Bail out if there's no stack map data.
-  assert((!CSInfos.empty() || (CSInfos.empty() && ConstPool.empty())) &&
+  assert((!CSInfos.empty() ||
+          (CSInfos.empty() && ConstPool.empty() && ConstSymPool.empty())) &&
          "Expected empty constant pool too!");
   assert((!CSInfos.empty() || (CSInfos.empty() && FnStackSize.empty())) &&
          "Expected empty function record too!");
   if (CSInfos.empty())
     return;
 
   MCContext &OutContext = AP.OutStreamer.getContext();
   MCStreamer &OS = AP.OutStreamer;
void StackMaps::serializeToStackMapSection() {
   emitFunctionFrameRecords(OS);
   emitConstantPoolEntries(OS);
   emitCallsiteEntries(OS);
   OS.AddBlankLine();
 
   // Clean up.
   CSInfos.clear();
   ConstPool.clear();
+  ConstSymPool.clear();
 }

test/CodeGen/X86/stackmap.ll

 ; RUN: llc < %s -mtriple=x86_64-apple-darwin -mcpu=corei7 | FileCheck %s
 ;
 ; Note: Print verbose stackmaps using -debug-only=stackmaps.
 
 ; CHECK-LABEL:  .section  __LLVM_STACKMAPS,__llvm_stackmaps
 ; CHECK-NEXT:  __LLVM_StackMaps:
 ; Header
 ; CHECK-NEXT:   .byte 1
 ; CHECK-NEXT:   .byte 0
 ; CHECK-NEXT:   .short 0
 ; Num Functions
 ; CHECK-NEXT:   .long 16
 ; Num LargeConstants
-; CHECK-NEXT:   .long 3
+; CHECK-NEXT:   .long 5
 ; Num Callsites
 ; CHECK-NEXT:   .long 20
 
 ; Functions and stack size
 ; CHECK-NEXT:   .quad _constantargs
 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _osrinline
 ; CHECK-NEXT:   .quad 24

 ; CHECK-NEXT:   .quad 56
 ; CHECK-NEXT:   .quad _needsStackRealignment
 ; CHECK-NEXT:   .quad -1
 
 ; Large Constants
 ; CHECK-NEXT:   .quad   2147483648
 ; CHECK-NEXT:   .quad   4294967295
 ; CHECK-NEXT:   .quad   4294967296
+; CHECK-NEXT:   .quad   _constSym1
+; CHECK-NEXT:   .quad   _constSym2
 
 ; Callsites
 ; Constant arguments
 ;
 ; CHECK-NEXT:   .quad   1
 ; CHECK-NEXT:   .long   L{{.*}}-_constantargs
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .short  12
+; CHECK-NEXT:   .short  15
 ; SmallConstant
 ; CHECK-NEXT:   .byte   4
 ; CHECK-NEXT:   .byte   8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .long   -1
 ; SmallConstant
 ; CHECK-NEXT:   .byte   4
 ; CHECK-NEXT:   .byte   8

 ; CHECK-NEXT:   .byte   8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .long   1
 ; LargeConstant at index 2
 ; CHECK-NEXT:   .byte   5
 ; CHECK-NEXT:   .byte   8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .long   2
+; LargeConstant at index 3
+; CHECK-NEXT:   .byte   5
+; CHECK-NEXT:   .byte   8
+; CHECK-NEXT:   .short  0
+; CHECK-NEXT:   .long   3
+; LargeConstant at index 3
+; CHECK-NEXT:   .byte   5
+; CHECK-NEXT:   .byte   8
+; CHECK-NEXT:   .short  0
+; CHECK-NEXT:   .long   3
+; LargeConstant at index 4
+; CHECK-NEXT:   .byte   5
+; CHECK-NEXT:   .byte   8
+; CHECK-NEXT:   .short  0
+; CHECK-NEXT:   .long   4
 ; SmallConstant
 ; CHECK-NEXT:   .byte   4
 ; CHECK-NEXT:   .byte   8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .long   -1
 
+
+@constSym1 = external constant i64
+@constSym2 = external constant i64
 define void @constantargs() {
 entry:
   %0 = inttoptr i64 12345 to i8*
-  tail call void (i64, i32, i8*, i32, ...)* @llvm.experimental.patchpoint.void(i64 1, i32 15, i8* %0, i32 0, i16 65535, i16 -1, i32 65536, i32 2000000000, i32 2147483647, i32 -1, i32 4294967295, i32 4294967296, i64 2147483648, i64 4294967295, i64 4294967296, i64 -1)
+  tail call void (i64, i32, i8*, i32, ...)* @llvm.experimental.patchpoint.void(i64 1, i32 15, i8* %0, i32 0, i16 65535, i16 -1, i32 65536, i32 2000000000, i32 2147483647, i32 -1, i32 4294967295, i32 4294967296, i64 2147483648, i64 4294967295, i64 4294967296, i64* @constSym1, i64* @constSym1, i64* @constSym2, i64 -1)
   ret void
 }
 
 ; Inline OSR Exit
 ;
 ; CHECK-LABEL:  .long   L{{.*}}-_osrinline
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  2