[Statepoints] Support for compressed pointers in the statepoint emission pass

[loicottet]

Summary

Adds a new GCStrategy to support two types of tracked pointers in the RewriteStatepointsForGC pass. When using this strategy, pointers in addrspace(2) will be saved as depot parameters to the rewritten statepoint calls, in addition to being inserted as GC parameters. This patch has been used by the LLVM backend of GraalVM's Native Image project in production for over a year with no major issues.

[apilipenko]

This seems oddly specific to the runtime you are working on. Why do you need to put compressed pointers into deopt bundes? Is it just to communicate the set of live compressed pointers to the runtime? In this case, coupling with deopt operand bundles seems incidental. We might be missing a more general mechanism for keeping track of different kinds of live GC pointers separately.

The information about the address space is preserved in the existing gc-live operand bundle. Can you use the address space of the pointers in the existing operand bundle for your purposes?

[loicottet]

In D150936#4363043, @apilipenko wrote:

This seems oddly specific to the runtime you are working on. Why do you need to put compressed pointers into deopt bundes? Is it just to communicate the set of live compressed pointers to the runtime? In this case, coupling with deopt operand bundles seems incidental. We might be missing a more general mechanism for keeping track of different kinds of live GC pointers separately.

I agree that a better mechanism to keep track of address spaces would be a better way to solve this.

The information about the address space is preserved in the existing gc-live operand bundle. Can you use the address space of the pointers in the existing operand bundle for your purposes?

The problem we are facing is that we need to retrieve this information from the stack maps, which do not preserve address space information. I don't know how doable it is to change the stack map format to accommodate this.

[apilipenko]

Right, one of the options would be to encode the information about GC pointer address spaces explicitly in the statepoint stack map.

Alternatively, you can have a downstream pass that rewrites the IR after RS4GC to accommodate the needs of your runtime. It would be a simple post-processing pass to parse gc-live operand bundle and add a deopt bundle with compressed pointers.

llvm/include/llvm/IR/GCStrategy.h

bool UseStatepoints = false; /// Uses gc.statepoints as opposed to gc.roots,
                                /// if set, NeededSafePoints and UsesMetadata
                                /// should be left at their default values.
 
   bool UseRS4GC = false; /// If UseStatepoints is set, this determines whether
                          /// the RewriteStatepointsForGC pass should rewrite
                          /// this function's calls.
                          /// This should only be set if UseStatepoints is set.
 
+  bool HasCompressedPointers =
+      false; /// Handles live pointers in addrspace(2)
+             /// as deopt values.
+             /// This should only be set if UseRS4GC is set.
+
   bool NeededSafePoints = false;    ///< if set, calls are inferred to be safepoints
   bool UsesMetadata = false;     ///< If set, backend must emit metadata tables.
 
 public:
   GCStrategy();
   virtual ~GCStrategy() = default;
 
   /// Return the name of the GC strategy.  This is the value of the collector
public:
   /// Returns true if the RewriteStatepointsForGC pass should run on functions
   /// using this GC.
   bool useRS4GC() const {
     assert((!UseRS4GC || useStatepoints()) &&
            "GC strategy has useRS4GC but not useStatepoints set");
     return UseRS4GC;
   }
 
+  /// Returns true if the RewriteStatepointsForGC pass should track pointers
+  /// in addrspace(2) as deopt parameters.
+  bool hasCompressedPointers() const {
+    assert((!HasCompressedPointers || useRS4GC()) &&
+           "GC strategy has hasCompressedPointers but not use RS4GC set");
+    return HasCompressedPointers;
+  }
+
   ///@}
 
   /// If set, appropriate metadata tables must be emitted by the back-end
   /// (assembler, JIT, or otherwise). The default stackmap information can be
   /// found in the StackMap section as described in the documentation.
   bool usesMetadata() const { return UsesMetadata; }
 
   /** @name GCRoot Specific Properties

llvm/include/llvm/IR/IRBuilder.h

public:
   CallInst *CreateGCStatepointCall(uint64_t ID, uint32_t NumPatchBytes,
                                    FunctionCallee ActualCallee, uint32_t Flags,
                                    ArrayRef<Value *> CallArgs,
                                    std::optional<ArrayRef<Use>> TransitionArgs,
                                    std::optional<ArrayRef<Use>> DeoptArgs,
                                    ArrayRef<Value *> GCArgs,
                                    const Twine &Name = "");
 
+  /// Create a call to the experimental.gc.statepoint intrinsic to
+  /// start a new statepoint sequence.
+  CallInst *CreateGCStatepointCall(uint64_t ID, uint32_t NumPatchBytes,
+                                   FunctionCallee ActualCallee, uint32_t Flags,
+                                   ArrayRef<Value *> CallArgs,
+                                   Optional<ArrayRef<Use>> TransitionArgs,
+                                   Optional<ArrayRef<Value *>> DeoptArgs,
+                                   ArrayRef<Value *> GCArgs,
+                                   const Twine &Name = "");
+
   /// Conveninence function for the common case when CallArgs are filled
   /// in using ArrayRef(CS.arg_begin(), CS.arg_end()); Use needs to be
   /// .get()'ed to get the Value pointer.
   CallInst *CreateGCStatepointCall(uint64_t ID, uint32_t NumPatchBytes,
                                    FunctionCallee ActualCallee,
                                    ArrayRef<Use> CallArgs,
                                    std::optional<ArrayRef<Value *>> DeoptArgs,
                                    ArrayRef<Value *> GCArgs,
public:
   /// start a new statepoint sequence.
   InvokeInst *CreateGCStatepointInvoke(
       uint64_t ID, uint32_t NumPatchBytes, FunctionCallee ActualInvokee,
       BasicBlock *NormalDest, BasicBlock *UnwindDest, uint32_t Flags,
       ArrayRef<Value *> InvokeArgs, std::optional<ArrayRef<Use>> TransitionArgs,
       std::optional<ArrayRef<Use>> DeoptArgs, ArrayRef<Value *> GCArgs,
       const Twine &Name = "");
 
+  /// Create an invoke to the experimental.gc.statepoint intrinsic to
+  /// start a new statepoint sequence.
+  InvokeInst *CreateGCStatepointInvoke(
+      uint64_t ID, uint32_t NumPatchBytes, FunctionCallee ActualInvokee,
+      BasicBlock *NormalDest, BasicBlock *UnwindDest, uint32_t Flags,
+      ArrayRef<Value *> InvokeArgs, Optional<ArrayRef<Use>> TransitionArgs,
+      Optional<ArrayRef<Value *>> DeoptArgs, ArrayRef<Value *> GCArgs,
+      const Twine &Name = "");
+
   // Convenience function for the common case when CallArgs are filled in using
   // ArrayRef(CS.arg_begin(), CS.arg_end()); Use needs to be .get()'ed to
   // get the Value *.
   InvokeInst *
   CreateGCStatepointInvoke(uint64_t ID, uint32_t NumPatchBytes,
                            FunctionCallee ActualInvokee, BasicBlock *NormalDest,
                            BasicBlock *UnwindDest, ArrayRef<Use> InvokeArgs,
                            std::optional<ArrayRef<Value *>> DeoptArgs,

llvm/lib/IR/BuiltinGCs.cpp

public:
   std::optional<bool> isGCManagedPointer(const Type *Ty) const override {
     // Method is only valid on pointer typed values.
     const PointerType *PT = cast<PointerType>(Ty);
     // We pick addrspace(1) as our GC managed heap.
     return (1 == PT->getAddressSpace());
   }
 };
 
+/// A GC strategy for the LLVM backend of the GraalVM compiler. The strategy
+/// is similar to the statepoint-example GC, but adds support for a second
+/// type of pointer in addrspace(2). Live pointers in this address space
+/// are inserted as deopt parameters as well as GC parameters in the
+/// statepoint intrinsic. As a consequence, this GC doesn't support regular
+/// deopt parameters
+class CompressedPointerGC : public GCStrategy {
+public:
+  CompressedPointerGC() {
+    UseStatepoints = true;
+    UseRS4GC = true;
+    HasCompressedPointers = true;
+    // These options are all gc.root specific, we specify them so that the
+    // gc.root lowering code doesn't run.
+    NeededSafePoints = false;
+    UsesMetadata = false;
+  }
+
+  Optional<bool> isGCManagedPointer(const Type *Ty) const override {
+    // Method is only valid on pointer typed values.
+    const PointerType *PT = cast<PointerType>(Ty);
+    // addrspace(1) represents absolute tracked pointers, and addrspace(2)
+    // represents heap base-relative pointers.
+    return (1 == PT->getAddressSpace() || 2 == PT->getAddressSpace());
+  }
+};
+
 } // end anonymous namespace
 
 // Register all the above so that they can be found at runtime.  Note that
 // these static initializers are important since the registration list is
 // constructed from their storage.
 static GCRegistry::Add<ErlangGC> A("erlang",
                                    "erlang-compatible garbage collector");
 static GCRegistry::Add<OcamlGC> B("ocaml", "ocaml 3.10-compatible GC");
 static GCRegistry::Add<ShadowStackGC>
     C("shadow-stack", "Very portable GC for uncooperative code generators");
 static GCRegistry::Add<StatepointGC> D("statepoint-example",
                                        "an example strategy for statepoint");
 static GCRegistry::Add<CoreCLRGC> E("coreclr", "CoreCLR-compatible GC");
+static GCRegistry::Add<CompressedPointerGC>
+    F("compressed-pointer", "GC supporting compressed pointers");
 
 // Provide hook to ensure the containing library is fully loaded.
 void llvm::linkAllBuiltinGCs() {}

llvm/lib/IR/IRBuilder.cpp

CallInst *IRBuilderBase::CreateGCStatepointCall(
     const Twine &Name) {
   return CreateGCStatepointCallCommon<Value *, Use, Use, Value *>(
       this, ID, NumPatchBytes, ActualCallee, Flags, CallArgs, TransitionArgs,
       DeoptArgs, GCArgs, Name);
 }
 
 CallInst *IRBuilderBase::CreateGCStatepointCall(
     uint64_t ID, uint32_t NumPatchBytes, FunctionCallee ActualCallee,
+    uint32_t Flags, ArrayRef<Value *> CallArgs,
+    Optional<ArrayRef<Use>> TransitionArgs,
+    Optional<ArrayRef<Value *>> DeoptArgs, ArrayRef<Value *> GCArgs,
+    const Twine &Name) {
+  return CreateGCStatepointCallCommon<Value *, Use, Value *, Value *>(
+      this, ID, NumPatchBytes, ActualCallee, Flags, CallArgs, TransitionArgs,
+      DeoptArgs, GCArgs, Name);
+}
+
+CallInst *IRBuilderBase::CreateGCStatepointCall(
+    uint64_t ID, uint32_t NumPatchBytes, FunctionCallee ActualCallee,
     ArrayRef<Use> CallArgs, std::optional<ArrayRef<Value *>> DeoptArgs,
     ArrayRef<Value *> GCArgs, const Twine &Name) {
   return CreateGCStatepointCallCommon<Use, Value *, Value *, Value *>(
       this, ID, NumPatchBytes, ActualCallee, uint32_t(StatepointFlags::None),
       CallArgs, std::nullopt, DeoptArgs, GCArgs, Name);
 }
 
 template <typename T0, typename T1, typename T2, typename T3>
InvokeInst *IRBuilderBase::CreateGCStatepointInvoke(
     const Twine &Name) {
   return CreateGCStatepointInvokeCommon<Value *, Use, Use, Value *>(
       this, ID, NumPatchBytes, ActualInvokee, NormalDest, UnwindDest, Flags,
       InvokeArgs, TransitionArgs, DeoptArgs, GCArgs, Name);
 }
 
 InvokeInst *IRBuilderBase::CreateGCStatepointInvoke(
     uint64_t ID, uint32_t NumPatchBytes, FunctionCallee ActualInvokee,
+    BasicBlock *NormalDest, BasicBlock *UnwindDest, uint32_t Flags,
+    ArrayRef<Value *> InvokeArgs, Optional<ArrayRef<Use>> TransitionArgs,
+    Optional<ArrayRef<Value *>> DeoptArgs, ArrayRef<Value *> GCArgs,
+    const Twine &Name) {
+  return CreateGCStatepointInvokeCommon<Value *, Use, Value *, Value *>(
+      this, ID, NumPatchBytes, ActualInvokee, NormalDest, UnwindDest, Flags,
+      InvokeArgs, TransitionArgs, DeoptArgs, GCArgs, Name);
+}
+
+InvokeInst *IRBuilderBase::CreateGCStatepointInvoke(
+    uint64_t ID, uint32_t NumPatchBytes, FunctionCallee ActualInvokee,
     BasicBlock *NormalDest, BasicBlock *UnwindDest, ArrayRef<Use> InvokeArgs,
     std::optional<ArrayRef<Value *>> DeoptArgs, ArrayRef<Value *> GCArgs,
     const Twine &Name) {
   return CreateGCStatepointInvokeCommon<Use, Value *, Value *, Value *>(
       this, ID, NumPatchBytes, ActualInvokee, NormalDest, UnwindDest,
       uint32_t(StatepointFlags::None), InvokeArgs, std::nullopt, DeoptArgs,
       GCArgs, Name);
 }

llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

if (CSAS.hasFnAttr(DeoptLowering))
       return CSAS.getFnAttr(DeoptLowering).getValueAsString();
     Function *F = Call->getCalledFunction();
     assert(F && F->hasFnAttribute(DeoptLowering));
     return F->getFnAttribute(DeoptLowering).getValueAsString();
   }
   return "live-through";
 }
 
+static void getCompressedGCArgs(std::vector<Value *> &CompressedArgsVector,
+                                ArrayRef<Value *> GCArgs) {
+  for (Value *LiveVal : GCArgs) {
+    if (auto *PT = dyn_cast<PointerType>(LiveVal->getType()))
+      if (PT->getAddressSpace() == 2)
+        CompressedArgsVector.push_back(LiveVal);
+  }
+}
+
 static void
 makeStatepointExplicitImpl(CallBase *Call, /* to replace */
                            const SmallVectorImpl<Value *> &BasePtrs,
                            const SmallVectorImpl<Value *> &LiveVariables,
                            PartiallyConstructedSafepointRecord &Result,
                            std::vector<DeferredReplacement> &Replacements,
                            const PointerToBaseTy &PointerToBase,
                            GCStrategy *GC) {
if (IID == Intrinsic::experimental_deoptimize) {
           F->getParent()
               ->getOrInsertFunction(GetFunctionName(IID, ElementSizeCI), FTy);
     }
   }
 
   // Create the statepoint given all the arguments
   GCStatepointInst *Token = nullptr;
   if (auto *CI = dyn_cast<CallInst>(Call)) {
-    CallInst *SPCall = Builder.CreateGCStatepointCall(
+    CallInst *SPCall;
+    if (GC->hasCompressedPointers()) {
+      assert(DeoptArgs->empty() &&
+             "Deopt args are not supported when using compressed pointers");
+      std::vector<Value *> CompressedArgsVector;
+      getCompressedGCArgs(CompressedArgsVector, GCArgs);
+      ArrayRef<Value *> CompressedArgs(CompressedArgsVector);
+
+      SPCall = Builder.CreateGCStatepointCall(
+          StatepointID, NumPatchBytes, CallTarget, Flags, CallArgs,
+          TransitionArgs, CompressedArgs, GCArgs, "safepoint_token");
+    } else {
+      SPCall = Builder.CreateGCStatepointCall(
           StatepointID, NumPatchBytes, CallTarget, Flags, CallArgs,
           TransitionArgs, DeoptArgs, GCArgs, "safepoint_token");
+    }
 
     SPCall->setTailCallKind(CI->getTailCallKind());
     SPCall->setCallingConv(CI->getCallingConv());
 
     // Currently we will fail on parameter attributes and on certain
     // function attributes.  In case if we can handle this set of attributes -
     // set up function attrs directly on statepoint and return attrs later for
     // gc_result intrinsic.
     SPCall->setAttributes(legalizeCallAttributes(
         CI->getContext(), CI->getAttributes(), SPCall->getAttributes()));
 
     Token = cast<GCStatepointInst>(SPCall);
 
     // Put the following gc_result and gc_relocate calls immediately after the
     // the old call (which we're about to delete)
     assert(CI->getNextNode() && "Not a terminator, must have next!");
     Builder.SetInsertPoint(CI->getNextNode());
     Builder.SetCurrentDebugLocation(CI->getNextNode()->getDebugLoc());
   } else {
     auto *II = cast<InvokeInst>(Call);
 
     // Insert the new invoke into the old block.  We'll remove the old one in a
     // moment at which point this will become the new terminator for the
     // original block.
-    InvokeInst *SPInvoke = Builder.CreateGCStatepointInvoke(
+    InvokeInst *SPInvoke;
+    if (GC->hasCompressedPointers()) {
+      assert(DeoptArgs->empty() &&
+             "Deopt args are not supported when using compressed pointers");
+      std::vector<Value *> CompressedArgsVector;
+      getCompressedGCArgs(CompressedArgsVector, GCArgs);
+      ArrayRef<Value *> CompressedArgs(CompressedArgsVector);
+
+      SPInvoke = Builder.CreateGCStatepointInvoke(
           StatepointID, NumPatchBytes, CallTarget, II->getNormalDest(),
-        II->getUnwindDest(), Flags, CallArgs, TransitionArgs, DeoptArgs, GCArgs,
-        "statepoint_token");
+          II->getUnwindDest(), Flags, CallArgs, TransitionArgs, CompressedArgs,
+          GCArgs, "statepoint_token");
+    } else {
+      SPInvoke = Builder.CreateGCStatepointInvoke(
+          StatepointID, NumPatchBytes, CallTarget, II->getNormalDest(),
+          II->getUnwindDest(), Flags, CallArgs, TransitionArgs, DeoptArgs,
+          GCArgs, "statepoint_token");
+    }
 
     SPInvoke->setCallingConv(II->getCallingConv());
 
     // Currently we will fail on parameter attributes and on certain
     // function attributes.  In case if we can handle this set of attributes -
     // set up function attrs directly on statepoint and return attrs later for
     // gc_result intrinsic.
     SPInvoke->setAttributes(legalizeCallAttributes(