Add GC preserve hook (mmtk#70)

This PR ports mmtk#58 to `dev`. This PR is mostly the same as mmtk#58 except that 1. this PR does not remove transitive pinning of shadow stack roots (we know it is unsound to remove the transitive pinning at this stage), and 2. this PR includes minor refactoring for GC codegen interface.

src/gc-interface.h

@@ -109,6 +109,13 @@ JL_DLLEXPORT const char* jl_gc_active_impl(void);
 // It still needs to be annotated with JL_DLLEXPORT since it is called from Rust by MMTk.
 JL_DLLEXPORT void jl_gc_sweep_stack_pools_and_mtarraylist_buffers(jl_ptls_t ptls) JL_NOTSAFEPOINT;
 
+// TODO: The preserve hook functions may be temporary. We should see the performance impact of the change.
+
+// Runtime hook for gc preserve begin. The GC needs to make sure that the preserved objects and its children stay alive and won't move.
+JL_DLLEXPORT void jl_gc_preserve_begin_hook(int n, ...) JL_NOTSAFEPOINT;
+// Runtime hook for gc preserve end. The GC needs to make sure that the preserved objects and its children stay alive and won't move.
+JL_DLLEXPORT void jl_gc_preserve_end_hook(void) JL_NOTSAFEPOINT;
+
 // ========================================================================= //
 // Metrics
 // ========================================================================= //

src/gc-mmtk.c

@@ -1208,6 +1208,53 @@ JL_DLLEXPORT jl_value_t *jl_gc_internal_obj_base_ptr(void *p)
     return NULL;
 }
 
+#define jl_p_gcpreserve_stack (jl_current_task->gcpreserve_stack)
+
+// This macro currently uses malloc instead of alloca because this function will exit
+// after pushing the roots into the gc_preserve_stack, which means that the preserve_begin function's
+// stack frame will be destroyed (together with its alloca variables). When we support lowering this code
+// inside the same function that is doing the preserve_begin/preserve_end calls we should be able to simple use allocas.
+// Note also that we use a separate stack for gc preserve roots to avoid the possibility of calling free
+// on a stack that has been allocated with alloca instead of malloc, which could happen depending on the order in which
+// JL_GC_POP() and jl_gc_preserve_end_hook() occurs.
+
+#define JL_GC_PUSHARGS_PRESERVE_ROOT_OBJS(rts_var,n)                                                    \
+  rts_var = ((jl_value_t**)malloc(((n)+2)*sizeof(jl_value_t*)))+2;                                      \
+  ((void**)rts_var)[-2] = (void*)JL_GC_ENCODE_PUSHARGS(n);                                              \
+  ((void**)rts_var)[-1] = jl_p_gcpreserve_stack;                                                        \
+  memset((void*)rts_var, 0, (n)*sizeof(jl_value_t*));                                                   \
+  jl_p_gcpreserve_stack = (jl_gcframe_t*)&(((void**)rts_var)[-2]);                                      \
+
+#define JL_GC_POP_PRESERVE_ROOT_OBJS()                                                                  \
+    jl_gcframe_t *curr = jl_p_gcpreserve_stack;                                                         \
+    if(curr) {                                                                                          \
+        (jl_p_gcpreserve_stack = jl_p_gcpreserve_stack->prev);                                          \
+        free(curr);                                                                                     \
+    }
+
+// Add each argument as a tpin root object.
+// However, we cannot use JL_GC_PUSH and JL_GC_POP since the slots should live
+// beyond this function. Instead, we maintain a tpin stack by mallocing/freeing
+// the frames for each of the preserve regions we encounter
+JL_DLLEXPORT void jl_gc_preserve_begin_hook(int n, ...) JL_NOTSAFEPOINT
+{
+    jl_value_t** frame;
+    JL_GC_PUSHARGS_PRESERVE_ROOT_OBJS(frame, n);
+    if (n == 0) return;
+
+    va_list args;
+    va_start(args, n);
+    for (int i = 0; i < n; i++) {
+        frame[i] = va_arg(args, jl_value_t *);
+    }
+    va_end(args);
+}
+
+JL_DLLEXPORT void jl_gc_preserve_end_hook(void) JL_NOTSAFEPOINT
+{
+    JL_GC_POP_PRESERVE_ROOT_OBJS();
+}
+
 #ifdef __cplusplus
 }
 #endif

src/gc-stock.c

@@ -4080,6 +4080,21 @@ JL_DLLEXPORT const char* jl_gc_active_impl(void) {
     return "Built with stock GC";
 }
 
+void jl_gc_notify_image_alloc(const char* img_data, size_t len)
+{
+    // Do nothing
+}
+
+JL_DLLEXPORT void jl_gc_preserve_begin_hook(int n, ...) JL_NOTSAFEPOINT
+{
+    jl_unreachable();
+}
+
+JL_DLLEXPORT void jl_gc_preserve_end_hook(void) JL_NOTSAFEPOINT
+{
+    jl_unreachable();
+}
+
 #ifdef __cplusplus
 }
 #endif

src/genericmemory.c

@@ -166,7 +166,7 @@ JL_DLLEXPORT jl_genericmemory_t *jl_ptr_to_genericmemory(jl_value_t *mtype, void
     m->length = nel;
     jl_genericmemory_data_owner_field(m) = own_buffer ? (jl_value_t*)m : NULL;
     if (own_buffer) {
-        PTR_PIN(m);
+        // FIXME: PTR_PIN(m);
         int isaligned = 0;  // TODO: allow passing memalign'd buffers
         jl_gc_track_malloced_genericmemory(ct->ptls, m, isaligned);
         size_t allocated_bytes = memory_block_usable_size(data, isaligned);

src/jl_exported_funcs.inc

@@ -181,6 +181,8 @@
     XX(jl_gc_set_max_memory) \
     XX(jl_gc_sync_total_bytes) \
     XX(jl_gc_total_hrtime) \
+    XX(jl_gc_preserve_begin_hook) \
+    XX(jl_gc_preserve_end_hook) \
     XX(jl_gdblookup) \
     XX(jl_generating_output) \
     XX(jl_declare_const_gf) \

src/julia_threads.h

@@ -274,6 +274,9 @@ typedef struct _jl_task_t {
     // uint48_t padding2_64;
     // saved gc stack top for context switches
     jl_gcframe_t *gcstack;
+    // GC stack of objects from gc preserve regions
+    // These must always be transitively pinned. Only used by MMTK.
+    jl_gcframe_t *gcpreserve_stack;
     size_t world_age;
     // quick lookup for current ptls
     jl_ptls_t ptls; // == jl_all_tls_states[tid]

src/llvm-final-gc-lowering.cpp

@@ -42,6 +42,8 @@ void FinalLowerGC::lowerPushGCFrame(CallInst *target, Function &F)
 
     IRBuilder<> builder(target);
     StoreInst *inst = builder.CreateAlignedStore(
+                // FIXME: We should use JL_GC_ENCODE_PUSHARGS_NO_TPIN here.
+                // We need to make sure things are properly pinned before turning this into a non TPIN push.
                 ConstantInt::get(T_size, JL_GC_ENCODE_PUSHARGS(nRoots)),
                 builder.CreateConstInBoundsGEP1_32(T_prjlvalue, gcframe, 0, "frame.nroots"),// GEP of 0 becomes a noop and eats the name
                 Align(sizeof(void*)));

src/llvm-gc-interface-passes.h

@@ -361,6 +361,7 @@ struct LateLowerGCFrame:  private JuliaPassContext {
     void PlaceGCFrameReset(State &S, unsigned R, unsigned MinColorRoot, ArrayRef<int> Colors, Value *GCFrame, Instruction *InsertBefore);
     void PlaceRootsAndUpdateCalls(ArrayRef<int> Colors, int PreAssignedColors, State &S, std::map<Value *, std::pair<int, int>>);
     void CleanupWriteBarriers(Function &F, State *S, const SmallVector<CallInst*, 0> &WriteBarriers, bool *CFGModified);
+    void CleanupGCPreserve(Function &F, CallInst *CI, Value *callee, Type *T_size);
     bool CleanupIR(Function &F, State *S, bool *CFGModified);
     void NoteUseChain(State &S, BBState &BBS, User *TheUser);
     SmallVector<int, 1> GetPHIRefinements(PHINode *phi, State &S);
@@ -413,4 +414,12 @@ struct FinalLowerGC: private JuliaPassContext {
     void lowerSafepoint(CallInst *target, Function &F);
 };
 
+inline bool isSpecialPtr(Type *Ty) {
+    PointerType *PTy = dyn_cast<PointerType>(Ty);
+    if (!PTy)
+        return false;
+    unsigned AS = PTy->getAddressSpace();
+    return AddressSpace::FirstSpecial <= AS && AS <= AddressSpace::LastSpecial;
+}
+
 #endif // LLVM_GC_PASSES_H

src/llvm-late-gc-lowering-mmtk.cpp

@@ -94,3 +94,49 @@ Value* LateLowerGCFrame::lowerGCAllocBytesLate(CallInst *target, Function &F)
     }
     return target;
 }
+
+void LateLowerGCFrame::CleanupGCPreserve(Function &F, CallInst *CI, Value *callee, Type *T_size) {
+    if (callee == gc_preserve_begin_func) {
+        // Initialize an IR builder.
+        IRBuilder<> builder(CI);
+
+        builder.SetCurrentDebugLocation(CI->getDebugLoc());
+        size_t nargs = 0;
+        State S2(F);
+
+        std::vector<Value*> args;
+        for (Use &U : CI->args()) {
+            Value *V = U;
+            if (isa<Constant>(V))
+                continue;
+            if (isa<PointerType>(V->getType())) {
+                if (isSpecialPtr(V->getType())) {
+                    int Num = Number(S2, V);
+                    if (Num >= 0) {
+                        nargs++;
+                        Value *Val = GetPtrForNumber(S2, Num, CI);
+                        args.push_back(Val);
+                    }
+                }
+            } else {
+                auto Nums = NumberAll(S2, V);
+                for (int Num : Nums) {
+                    if (Num < 0)
+                        continue;
+                    Value *Val = GetPtrForNumber(S2, Num, CI);
+                    args.push_back(Val);
+                    nargs++;
+                }
+            }
+        }
+        args.insert(args.begin(), ConstantInt::get(T_size, nargs));
+
+        ArrayRef<Value*> args_llvm = ArrayRef<Value*>(args);
+        builder.CreateCall(getOrDeclare(jl_well_known::GCPreserveBeginHook), args_llvm );
+    } else if (callee == gc_preserve_end_func) {
+        // Initialize an IR builder.
+        IRBuilder<> builder(CI);
+        builder.SetCurrentDebugLocation(CI->getDebugLoc());
+        builder.CreateCall(getOrDeclare(jl_well_known::GCPreserveEndHook), {});
+    }
+}

src/llvm-late-gc-lowering-stock.cpp

@@ -7,3 +7,7 @@ Value* LateLowerGCFrame::lowerGCAllocBytesLate(CallInst *target, Function &F)
     // Do nothing for the stock GC
     return target;
 }
+
+void LateLowerGCFrame::CleanupGCPreserve(Function &F, CallInst *CI, Value *callee, Type *T_size) {
+    // Do nothing for the stock GC
+}

src/llvm-late-gc-lowering.cpp

@@ -13,14 +13,6 @@ static bool isTrackedValue(Value *V) {
     return PT && PT->getAddressSpace() == AddressSpace::Tracked;
 }
 
-static bool isSpecialPtr(Type *Ty) {
-    PointerType *PTy = dyn_cast<PointerType>(Ty);
-    if (!PTy)
-        return false;
-    unsigned AS = PTy->getAddressSpace();
-    return AddressSpace::FirstSpecial <= AS && AS <= AddressSpace::LastSpecial;
-}
-
 // return how many Special pointers are in T (count > 0),
 // and if there is anything else in T (all == false)
 CountTrackedPointers::CountTrackedPointers(Type *T, bool ignore_loaded) {
@@ -2006,9 +1998,11 @@ bool LateLowerGCFrame::CleanupIR(Function &F, State *S, bool *CFGModified) {
                 continue;
             }
             Value *callee = CI->getCalledOperand();
-            if (callee && (callee == gc_flush_func || callee == gc_preserve_begin_func
-                        || callee == gc_preserve_end_func)) {
+            if (callee && callee == gc_flush_func) {
                 /* No replacement */
+            } else if (callee && (callee == gc_preserve_begin_func
+                        || callee == gc_preserve_end_func)) {
+                CleanupGCPreserve(F, CI, callee, T_size);
             } else if (pointer_from_objref_func != nullptr && callee == pointer_from_objref_func) {
                 auto *obj = CI->getOperand(0);
                 auto *ASCI = new AddrSpaceCastInst(obj, CI->getType(), "", CI);

src/llvm-pass-helpers.cpp

@@ -252,6 +252,8 @@ namespace jl_well_known {
     static const char *GC_SMALL_ALLOC_NAME = XSTR(jl_gc_small_alloc);
     static const char *GC_QUEUE_ROOT_NAME = XSTR(jl_gc_queue_root);
     static const char *GC_ALLOC_TYPED_NAME = XSTR(jl_gc_alloc_typed);
+    static const char *GC_PRESERVE_BEGIN_HOOK_NAME = XSTR(jl_gc_preserve_begin_hook);
+    static const char *GC_PRESERVE_END_HOOK_NAME = XSTR(jl_gc_preserve_end_hook);
 
     using jl_intrinsics::addGCAllocAttributes;
 
@@ -320,4 +322,50 @@ namespace jl_well_known {
             allocTypedFunc->addFnAttr(Attribute::getWithAllocSizeArgs(ctx, 1, None));
             return addGCAllocAttributes(allocTypedFunc);
         });
+
+    const WellKnownFunctionDescription GCPreserveBeginHook(
+        GC_PRESERVE_BEGIN_HOOK_NAME,
+        [](Type *T_size) {
+            auto &ctx = T_size->getContext();
+            auto func = Function::Create(
+                FunctionType::get(
+                    Type::getVoidTy(ctx),
+                    { T_size },
+                    true),
+                Function::ExternalLinkage,
+                GC_PRESERVE_BEGIN_HOOK_NAME);
+
+#if JL_LLVM_VERSION >= 160000
+            func->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
+#else
+            func->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
+#endif
+            return func;
+        });
+
+    const WellKnownFunctionDescription GCPreserveEndHook(
+        GC_PRESERVE_END_HOOK_NAME,
+        [](Type *T_size) {
+            auto &ctx = T_size->getContext();
+            auto func = Function::Create(
+                FunctionType::get(
+                    Type::getVoidTy(ctx),
+                    {  },
+                    false),
+                Function::ExternalLinkage,
+                GC_PRESERVE_END_HOOK_NAME);
+#if JL_LLVM_VERSION >= 160000
+            func->setMemoryEffects(MemoryEffects::inaccessibleOrArgMemOnly());
+#else
+            func->addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
+#endif
+            return func;
+        });
+}
+
+void setName(llvm::Value *V, const llvm::Twine &Name, int debug_info)
+{
+    if (debug_info >= 2 && !llvm::isa<llvm::Constant>(V)) {
+        V->setName(Name);
+    }
 }

src/llvm-pass-helpers.h

@@ -155,6 +155,12 @@ namespace jl_well_known {
 
     // `jl_gc_alloc_typed`: allocates bytes.
     extern const WellKnownFunctionDescription GCAllocTyped;
+
+    // `jl_gc_preserve_begin_hook`: called at the beginning of gc preserve regions, if required
+    extern const WellKnownFunctionDescription GCPreserveBeginHook;
+
+    // `jl_gc_preserve_end_hook`: called at the end of gc preserve regions, if required
+    extern const WellKnownFunctionDescription GCPreserveEndHook;
 }
 
 void setName(llvm::Value *V, const llvm::Twine &Name, int debug_info);