Add GC preserve hook (#70)

This PR ports #58 to `dev`. This PR is mostly the same as #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/Makefile

@@ -29,10 +29,6 @@ ifeq ($(USECLANG),1)
 FLAGS += -Wno-return-type-c-linkage -Wno-atomic-alignment
 endif
 
-ifeq ($(WITH_MMTK), 1)
-FLAGS += -I$(MMTK_API_INC)
-endif
-
 FLAGS += -DJL_BUILD_ARCH='"$(ARCH)"'
 ifeq ($(OS),WINNT)
 FLAGS += -DJL_BUILD_UNAME='"NT"'
@@ -65,7 +61,14 @@ RT_LLVMLINK :=
 CG_LLVMLINK :=
 
 ifeq ($(JULIACODEGEN),LLVM)
+# Currently these files are used by both GCs. But we should make the list specific to stock, and MMTk should have its own implementation.
 GC_CODEGEN_SRCS := llvm-final-gc-lowering llvm-late-gc-lowering llvm-gc-invariant-verifier
+ifeq ($(WITH_MMTK), 1)
+FLAGS += -I$(MMTK_API_INC)
+GC_CODEGEN_SRCS += llvm-late-gc-lowering-mmtk
+else
+GC_CODEGEN_SRCS += llvm-late-gc-lowering-stock
+endif
 CODEGEN_SRCS := codegen jitlayers aotcompile debuginfo disasm llvm-simdloop \
 	llvm-pass-helpers llvm-ptls \
 	llvm-lower-handlers llvm-propagate-addrspaces \

src/gc-interface.h

@@ -103,6 +103,13 @@ JL_DLLEXPORT unsigned char jl_gc_pin_object(void* obj);
 // Returns the version of which GC implementation is being used according to the list of supported GCs
 JL_DLLEXPORT const char* jl_active_gc_impl(void);
 
+// 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

@@ -1240,6 +1240,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

@@ -3999,6 +3999,16 @@ JL_DLLEXPORT void jl_gc_wb2_slow(const void *parent, const void* ptr) JL_NOTSAFE
 {
 }
 
+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();
+}
+
 JL_DLLEXPORT const char* jl_active_gc_impl(void) {
     return "";
 }

src/jl_exported_funcs.inc

@@ -186,6 +186,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.h

@@ -2367,6 +2367,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

@@ -360,6 +360,7 @@ struct LateLowerGCFrame:  private JuliaPassContext {
     void PlaceGCFrameStores(State &S, unsigned MinColorRoot, ArrayRef<int> Colors, Value *GCFrame);
     void PlaceRootsAndUpdateCalls(SmallVectorImpl<int> &Colors, 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);
@@ -427,4 +428,12 @@ struct FinalLowerGC: private JuliaPassContext {
 #endif
 };
 
+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

@@ -0,0 +1,47 @@
+#include "llvm-gc-interface-passes.h"
+
+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

@@ -0,0 +1,5 @@
+#include "llvm-gc-interface-passes.h"
+
+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) {
@@ -2089,9 +2081,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

@@ -347,6 +347,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);
 #ifdef MMTK_GC
     static const char *GC_WB_1_NAME = XSTR(jl_gc_wb1_noinline);
     static const char *GC_WB_2_NAME = XSTR(jl_gc_wb2_noinline);
@@ -426,6 +428,45 @@ namespace jl_well_known {
             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;
+        });
+
 #ifdef MMTK_GC
     const WellKnownFunctionDescription GCWriteBarrier1(
         GC_WB_1_NAME,

src/llvm-pass-helpers.h

@@ -163,6 +163,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;
+
 #ifdef MMTK_GC
     extern const WellKnownFunctionDescription GCWriteBarrier1;
     extern const WellKnownFunctionDescription GCWriteBarrier2;