Merge pull request #23610 from JuliaLang/kf/fancygcpreserve

Fancier GC preserve intrinsics

Author: Keno

base/boot.jl

@@ -438,6 +438,4 @@ show(@nospecialize a) = show(STDOUT, a)
 print(@nospecialize a...) = print(STDOUT, a...)
 println(@nospecialize a...) = println(STDOUT, a...)
 
-gcuse(@nospecialize a) = ccall(:jl_gc_use, Void, (Any,), a)
-
 ccall(:jl_set_istopmod, Void, (Any, Bool), Core, true)

base/pointer.jl

@@ -150,3 +150,24 @@ isless(x::Ptr, y::Ptr) = isless(UInt(x), UInt(y))
 +(x::Ptr, y::Integer) = oftype(x, (UInt(x) + (y % UInt) % UInt))
 -(x::Ptr, y::Integer) = oftype(x, (UInt(x) - (y % UInt) % UInt))
 +(x::Integer, y::Ptr) = y + x
+
+"""
+Temporarily protects an object from being garbage collected, even
+if it would otherwise be unreferenced.
+
+The last argument is the expression to preserve objects during.
+The previous arguments are the objects to preserve.
+"""
+macro gc_preserve(args...)
+    syms = args[1:end-1]
+    for x in syms
+        isa(x, Symbol) || error("Preserved variable must be a symbol")
+    end
+    s, r = gensym(), gensym()
+    esc(quote
+        $s = $(Expr(:gc_preserve_begin, syms...))
+        $r = $(args[end])
+        $(Expr(:gc_preserve_end, s))
+        $r
+    end)
+end

base/strings/string.jl

@@ -73,10 +73,7 @@ codeunit(s::AbstractString, i::Integer)
     @boundscheck if (i < 1) | (i > sizeof(s))
         throw(BoundsError(s,i))
     end
-    ptr = pointer(s, i)
-    r = unsafe_load(ptr)
-    Core.gcuse(s)
-    r
+    @gc_preserve s unsafe_load(pointer(s, i))
 end
 
 write(io::IO, s::String) = unsafe_write(io, pointer(s), reinterpret(UInt, sizeof(s)))

doc/src/devdocs/llvm.md

@@ -297,3 +297,24 @@ for the function. As a result, the external rooting must be arranged while the
 value is still tracked by the system. I.e. it is not valid to attempt to use the
 result of this operation to establish a global root - the optimizer may have
 already dropped the value.
+
+### Keeping values alive in the absence of uses
+
+In certain cases it is necessary to keep an object alive, even though there is
+no compiler-visible use of said object. This may be case for low level code
+that operates on the memory-representation of an object directly or code that
+needs to interface with C code. In order to allow this, we provide the following
+intrinsics at the LLVM level:
+```
+token @llvm.julia.gc_preserve_begin(...)
+void @llvm.julia.gc_preserve_end(token)
+```
+(The `llvm.` in the name is required in order to be able to use the `token`
+type). The semantics of these intrinsics are as follows:
+At any safepoint that is dominated by a `gc_preserve_begin` call, but that is not
+not dominated by a corresponding `gc_preserve_end` call (i.e. a call whose argument
+is the token returned by a `gc_preserve_begin` call), the values passed as
+arguments to that `gc_preserve_begin` will be kept live. Note that the
+`gc_preserve_begin` still counts as a regular use of those values, so the
+standard lifetime semantics will ensure that the values will be kept alive
+before entering the preserve region.

src/ast.c

@@ -60,6 +60,7 @@ jl_sym_t *isdefined_sym; jl_sym_t *nospecialize_sym;
 jl_sym_t *macrocall_sym;
 jl_sym_t *hygienicscope_sym;
 jl_sym_t *escape_sym;
+jl_sym_t *gc_preserve_begin_sym; jl_sym_t *gc_preserve_end_sym;
 
 static uint8_t flisp_system_image[] = {
 #include <julia_flisp.boot.inc>
@@ -340,6 +341,8 @@ void jl_init_frontend(void)
     macrocall_sym = jl_symbol("macrocall");
     escape_sym = jl_symbol("escape");
     hygienicscope_sym = jl_symbol("hygienic-scope");
+    gc_preserve_begin_sym = jl_symbol("gc_preserve_begin");
+    gc_preserve_end_sym = jl_symbol("gc_preserve_end");
 }
 
 JL_DLLEXPORT void jl_lisp_prompt(void)

src/ccall.cpp

@@ -1679,13 +1679,6 @@ static jl_cgval_t emit_ccall(jl_codectx_t &ctx, jl_value_t **args, size_t nargs)
         emit_signal_fence(ctx);
         return ghostValue(jl_void_type);
     }
-    else if (is_libjulia_func(jl_gc_use)) {
-        assert(lrt == T_void);
-        assert(!isVa && !llvmcall && nargt == 1);
-        ctx.builder.CreateCall(prepare_call(gc_use_func), {decay_derived(boxed(ctx, argv[0]))});
-        JL_GC_POP();
-        return ghostValue(jl_void_type);
-    }
     else if (_is_libjulia_func((uintptr_t)ptls_getter, "jl_get_ptls_states")) {
         assert(lrt == T_size);
         assert(!isVa && !llvmcall && nargt == 0);

src/codegen.cpp

@@ -356,7 +356,8 @@ static GlobalVariable *jlgetworld_global;
 
 // placeholder functions
 static Function *gcroot_flush_func;
-static Function *gc_use_func;
+static Function *gc_preserve_begin_func;
+static Function *gc_preserve_end_func;
 static Function *except_enter_func;
 static Function *pointer_from_objref_func;
 
@@ -3905,6 +3906,42 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr)
     else if (head == boundscheck_sym) {
         return mark_julia_const(bounds_check_enabled(ctx, jl_true) ? jl_true : jl_false);
     }
+    else if (head == gc_preserve_begin_sym) {
+        size_t nargs = jl_array_len(ex->args);
+        jl_cgval_t *argv = (jl_cgval_t*)alloca(sizeof(jl_cgval_t) * nargs);
+        for (size_t i = 0; i < nargs; ++i) {
+            argv[i] = emit_expr(ctx, args[i]);
+        }
+        size_t nargsboxed = 0;
+        Value **vals = (Value**)alloca(sizeof(Value *) * nargs);
+        for (size_t i = 0; i < nargs; ++i) {
+            if (!argv[i].isboxed) {
+                // This is intentionally not an error to allow writing
+                // generic code more easily.
+                continue;
+            } else if (argv[i].constant) {
+                continue;
+            }
+            vals[nargsboxed++] = argv[i].Vboxed;
+        }
+        Value *token = ctx.builder.CreateCall(prepare_call(gc_preserve_begin_func),
+            ArrayRef<Value*>(vals, nargsboxed));
+        jl_cgval_t tok(token, NULL, false, (jl_value_t*)jl_void_type, NULL);
+        tok.isimmutable = true;
+        return tok;
+    }
+    else if (head == gc_preserve_end_sym) {
+        // We only support ssa values as the argument. Everything else will
+        // fall back to the default behavior of preserving the argument value
+        // until the end of the scope, which is correct, but not optimal.
+        if (!jl_is_ssavalue(args[0])) {
+            return jl_cgval_t((jl_value_t*)jl_void_type);
+        }
+        jl_cgval_t token = emit_expr(ctx, args[0]);
+        assert(token.V->getType()->isTokenTy());
+        ctx.builder.CreateCall(prepare_call(gc_preserve_end_func), {token.V});
+        return jl_cgval_t((jl_value_t*)jl_void_type);
+    }
     else {
         if (!strcmp(jl_symbol_name(head), "$"))
             jl_error("syntax: prefix \"$\" in non-quoted expression");
@@ -6526,11 +6563,17 @@ static void init_julia_llvm_env(Module *m)
                                          "julia.gcroot_flush");
     add_named_global(gcroot_flush_func, (void*)NULL, /*dllimport*/false);
 
-    gc_use_func = Function::Create(FunctionType::get(T_void,
-                                         ArrayRef<Type*>(PointerType::get(T_jlvalue, AddressSpace::Derived)), false),
+    gc_preserve_begin_func = Function::Create(FunctionType::get(Type::getTokenTy(jl_LLVMContext),
+                                         ArrayRef<Type*>(), true),
                                          Function::ExternalLinkage,
-                                         "julia.gc_use");
-    add_named_global(gc_use_func, (void*)NULL, /*dllimport*/false);
+                                         "llvm.julia.gc_preserve_begin");
+    add_named_global(gc_preserve_begin_func, (void*)NULL, /*dllimport*/false);
+
+    gc_preserve_end_func = Function::Create(FunctionType::get(T_void,
+                                        ArrayRef<Type*>(Type::getTokenTy(jl_LLVMContext)), false),
+                                        Function::ExternalLinkage,
+                                        "llvm.julia.gc_preserve_end");
+    add_named_global(gc_preserve_end_func, (void*)NULL, /*dllimport*/false);
 
     pointer_from_objref_func = Function::Create(FunctionType::get(T_size,
                                          ArrayRef<Type*>(PointerType::get(T_jlvalue, AddressSpace::Derived)), false),

src/interpreter.c

@@ -506,6 +506,11 @@ static jl_value_t *eval(jl_value_t *e, interpreter_state *s)
     else if (ex->head == boundscheck_sym || ex->head == inbounds_sym || ex->head == fastmath_sym ||
              ex->head == simdloop_sym || ex->head == meta_sym) {
         return jl_nothing;
+    } else if (ex->head == gc_preserve_begin_sym || ex->head == gc_preserve_end_sym) {
+        // The interpreter generally keeps values that were assigned in this scope
+        // rooted. If the interpreter learns to be more agressive here, we may
+        // want to explicitly root these values.
+        return jl_nothing;
     }
     jl_errorf("unsupported or misplaced expression %s", jl_symbol_name(ex->head));
     abort();

src/julia-syntax.scm

@@ -3084,7 +3084,8 @@ f(x) = yt(x)
                                (memq (car e) '(quote top core line inert local local-def unnecessary
                                                meta inbounds boundscheck simdloop decl
                                                implicit-global global globalref outerref
-                                               const = null method call foreigncall ssavalue))))
+                                               const = null method call foreigncall ssavalue
+                                               gc_preserve_begin gc_preserve_end))))
                          (lam:body lam))))
                (unused (map cadr (filter (lambda (x) (memq (car x) '(method =)))
                                          leading))))
@@ -3816,8 +3817,13 @@ f(x) = yt(x)
              (if tail (emit-return '(null)))
              '(null))
 
+            ((gc_preserve_begin)
+              (let ((s (make-ssavalue)))
+                (emit `(= ,s ,e))
+                s))
+
             ;; other top level expressions and metadata
-            ((import importall using export line meta inbounds boundscheck simdloop)
+            ((import importall using export line meta inbounds boundscheck simdloop gc_preserve_end)
              (let ((have-ret? (and (pair? code) (pair? (car code)) (eq? (caar code) 'return))))
                (cond ((eq? (car e) 'line)
                       (set! current-loc e)

src/julia_internal.h

@@ -999,6 +999,7 @@ extern jl_sym_t *propagate_inbounds_sym;
 extern jl_sym_t *isdefined_sym;
 extern jl_sym_t *nospecialize_sym;
 extern jl_sym_t *boundscheck_sym;
+extern jl_sym_t *gc_preserve_begin_sym; extern jl_sym_t *gc_preserve_end_sym;
 
 void jl_register_fptrs(uint64_t sysimage_base, const char *base, const int32_t *offsets,
                        jl_method_instance_t **linfos, size_t n);

src/llvm-late-gc-lowering.cpp

@@ -259,6 +259,10 @@ struct State {
     // of the value (but not the other way around).
     std::map<int, int> LoadRefinements;
 
+    // GC preserves map. All safepoints dominated by the map key, but not any
+    // of its uses need to preserve the values listed in the map value.
+    std::map<Instruction *, std::vector<int>> GCPreserves;
+
     // The assignment of numbers to safepoints. The indices in the map
     // are indices into the next three maps which store safepoint properties
     std::map<Instruction *, int> SafepointNumbering;
@@ -322,7 +326,8 @@ struct LateLowerGCFrame: public FunctionPass {
     MDNode *tbaa_tag;
     Function *ptls_getter;
     Function *gc_flush_func;
-    Function *gc_use_func;
+    Function *gc_preserve_begin_func;
+    Function *gc_preserve_end_func;
     Function *pointer_from_objref_func;
     Function *alloc_obj_func;
     Function *pool_alloc_func;
@@ -754,17 +759,29 @@ State LateLowerGCFrame::LocalScan(Function &F) {
                 if (CI->canReturnTwice()) {
                     S.ReturnsTwice.push_back(CI);
                 }
-                if (isa<IntrinsicInst>(CI)) {
-                    // Intrinsics are never safepoints.
-                    continue;
-                }
                 if (auto callee = CI->getCalledFunction()) {
+                    if (callee == gc_preserve_begin_func) {
+                        std::vector<int> args;
+                        for (Use &U : CI->arg_operands()) {
+                            Value *V = U;
+                            int Num = Number(S, V);
+                            if (Num >= 0)
+                                args.push_back(Num);
+                        }
+                        S.GCPreserves[CI] = args;
+                        continue;
+                    }
                     // Known functions emitted in codegen that are not safepoints
-                    if (callee == pointer_from_objref_func || callee == gc_use_func ||
+                    if (callee == pointer_from_objref_func || callee == gc_preserve_begin_func ||
+                        callee == gc_preserve_end_func ||
                         callee->getName() == "memcmp") {
                         continue;
                     }
                 }
+                if (isa<IntrinsicInst>(CI)) {
+                    // Intrinsics are never safepoints.
+                    continue;
+                }
                 int SafepointNumber = NoteSafepoint(S, BBS, CI);
                 BBS.HasSafepoint = true;
                 BBS.TopmostSafepoint = SafepointNumber;
@@ -912,6 +929,7 @@ JL_USED_FUNC static void dumpSafepointsForBBName(Function &F, State &S, const ch
 }
 
 void LateLowerGCFrame::ComputeLiveSets(Function &F, State &S) {
+    DominatorTree *DT = nullptr;
     // Iterate over all safe points. Add to live sets all those variables that
     // are now live across their parent block.
     for (auto it : S.SafepointNumbering) {
@@ -935,6 +953,33 @@ void LateLowerGCFrame::ComputeLiveSets(Function &F, State &S) {
             if (RefinedPtr == -1 || HasBitSet(LS, RefinedPtr))
                 LS[Idx] = 0;
         }
+        // If the function has GC preserves, figure out whether we need to
+        // add in any extra live values.
+        if (!S.GCPreserves.empty()) {
+            if (!DT) {
+                DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+            }
+            for (auto it2 : S.GCPreserves) {
+                if (!DT->dominates(it2.first, Safepoint))
+                    continue;
+                bool OutsideRange = false;
+                for (const User *U : it2.first->users()) {
+                    // If this is dominated by an end, we don't need to add
+                    // the values to our live set.
+                    if (DT->dominates(cast<Instruction>(U), Safepoint)) {
+                        OutsideRange = true;
+                        break;
+                    }
+                }
+                if (OutsideRange)
+                    continue;
+                for (unsigned Num : it2.second) {
+                    if (Num >= LS.size())
+                        LS.resize(Num + 1);
+                    LS[Num] = 1;
+                }
+            }
+        }
     }
     // Compute the interference graph
     for (int i = 0; i <= S.MaxPtrNumber; ++i) {
@@ -1153,8 +1198,8 @@ bool LateLowerGCFrame::CleanupIR(Function &F) {
             }
             CallingConv::ID CC = CI->getCallingConv();
             auto callee = CI->getCalledValue();
-            if ((gc_flush_func != nullptr && callee == gc_flush_func) ||
-                (gc_use_func != nullptr && callee == gc_use_func)) {
+            if (callee && (callee == gc_flush_func || callee == gc_preserve_begin_func
+                        || callee == gc_preserve_end_func)) {
                 /* No replacement */
             } else if (pointer_from_objref_func != nullptr && callee == pointer_from_objref_func) {
                 auto *obj = CI->getOperand(0);
@@ -1244,6 +1289,9 @@ bool LateLowerGCFrame::CleanupIR(Function &F) {
                 NewCall->takeName(CI);
                 CI->replaceAllUsesWith(NewCall);
             }
+            if (!CI->use_empty()) {
+                CI->replaceAllUsesWith(UndefValue::get(CI->getType()));
+            }
             it = CI->eraseFromParent();
             ChangesMade = true;
         }
@@ -1423,7 +1471,8 @@ static void addRetNoAlias(Function *F)
 bool LateLowerGCFrame::DefineFunctions(Module &M) {
     ptls_getter = M.getFunction("jl_get_ptls_states");
     gc_flush_func = M.getFunction("julia.gcroot_flush");
-    gc_use_func = M.getFunction("julia.gc_use");
+    gc_preserve_begin_func = M.getFunction("llvm.julia.gc_preserve_begin");
+    gc_preserve_end_func = M.getFunction("llvm.julia.gc_preserve_end");
     pointer_from_objref_func = M.getFunction("julia.pointer_from_objref");
     auto &ctx = M.getContext();
     T_size = M.getDataLayout().getIntPtrType(ctx);

src/macroexpand.scm

@@ -307,7 +307,7 @@
                                    ,(resolve-expansion-vars-with-new-env (caddr arg) env m inarg))))
                              (else
                               `(global ,(resolve-expansion-vars-with-new-env arg env m inarg))))))
-           ((using import importall export meta line inbounds boundscheck simdloop) (map unescape e))
+           ((using import importall export meta line inbounds boundscheck simdloop gc_preserve gc_preserve_end) (map unescape e))
            ((macrocall) e) ; invalid syntax anyways, so just act like it's quoted.
            ((symboliclabel) e)
            ((symbolicgoto) e)

test/llvmpasses/gcroots.ll

@@ -223,3 +223,26 @@ top:
     ret void
 }
 
+declare token @llvm.julia.gc_preserve_begin(...)
+declare void @llvm.julia.gc_preserve_end(token)
+
+define void @gc_preserve(i64 %a) {
+; CHECK-LABEL: @gc_preserve
+; CHECK: %gcframe = alloca %jl_value_t addrspace(10)*, i32 4
+top:
+    %ptls = call %jl_value_t*** @jl_get_ptls_states()
+    %aboxed = call %jl_value_t addrspace(10)* @jl_box_int64(i64 signext %a)
+; CHECK: store %jl_value_t addrspace(10)* %aboxed
+    call void @jl_safepoint()
+    %tok = call token (...) @llvm.julia.gc_preserve_begin(%jl_value_t addrspace(10)* %aboxed)
+    %aboxed2 = call %jl_value_t addrspace(10)* @jl_box_int64(i64 signext %a)
+; CHECK: store %jl_value_t addrspace(10)* %aboxed2
+    call void @jl_safepoint()
+    call void @llvm.julia.gc_preserve_end(token %tok)
+    %aboxed3 = call %jl_value_t addrspace(10)* @jl_box_int64(i64 signext %a)
+; CHECK: store %jl_value_t addrspace(10)* %aboxed3
+    call void @jl_safepoint()
+    call void @one_arg_boxed(%jl_value_t addrspace(10)* %aboxed2)
+    call void @one_arg_boxed(%jl_value_t addrspace(10)* %aboxed3)
+    ret void
+}