Refetch PTLS via Task struct (preparation for task migration)

src/ccall.cpp

@@ -1458,7 +1458,7 @@ static jl_cgval_t emit_ccall(jl_codectx_t &ctx, jl_value_t **args, size_t nargs)
         JL_GC_POP();
         ctx.builder.CreateCall(prepare_call(gcroot_flush_func));
         emit_signal_fence(ctx);
-        ctx.builder.CreateLoad(T_size, ctx.signalPage, true);
+        ctx.builder.CreateLoad(T_size, get_current_signal_page(ctx), true);
         emit_signal_fence(ctx);
         return ghostValue(jl_nothing_type);
     }
@@ -1467,14 +1467,14 @@ static jl_cgval_t emit_ccall(jl_codectx_t &ctx, jl_value_t **args, size_t nargs)
         assert(!isVa && !llvmcall && nccallargs == 0);
         JL_GC_POP();
         return mark_or_box_ccall_result(ctx,
-            ctx.builder.CreatePtrToInt(ctx.ptlsStates, lrt),
+            ctx.builder.CreatePtrToInt(get_current_ptls(ctx), lrt),
             retboxed, rt, unionall, static_rt);
     }
     else if (is_libjulia_func(jl_threadid)) {
         assert(lrt == T_int16);
         assert(!isVa && !llvmcall && nccallargs == 0);
         JL_GC_POP();
-        Value *ptls_i16 = emit_bitcast(ctx, ctx.ptlsStates, T_pint16);
+        Value *ptls_i16 = emit_bitcast(ctx, get_current_ptls(ctx), T_pint16);
         const int tid_offset = offsetof(jl_tls_states_t, tid);
         Value *ptid = ctx.builder.CreateInBoundsGEP(ptls_i16, ConstantInt::get(T_size, tid_offset / 2));
         LoadInst *tid = ctx.builder.CreateAlignedLoad(ptid, Align(sizeof(int16_t)));
@@ -1485,18 +1485,14 @@ static jl_cgval_t emit_ccall(jl_codectx_t &ctx, jl_value_t **args, size_t nargs)
         assert(lrt == T_prjlvalue);
         assert(!isVa && !llvmcall && nccallargs == 0);
         JL_GC_POP();
-        Value *ptls_pv = emit_bitcast(ctx, ctx.ptlsStates, T_pprjlvalue);
-        const int ct_offset = offsetof(jl_tls_states_t, current_task);
-        Value *pct = ctx.builder.CreateInBoundsGEP(ptls_pv, ConstantInt::get(T_size, ct_offset / sizeof(void*)));
-        LoadInst *ct = ctx.builder.CreateAlignedLoad(pct, Align(sizeof(void*)));
-        tbaa_decorate(tbaa_const, ct);
+        auto ct = get_current_task(ctx);
         return mark_or_box_ccall_result(ctx, ct, retboxed, rt, unionall, static_rt);
     }
     else if (is_libjulia_func(jl_set_next_task)) {
         assert(lrt == T_void);
         assert(!isVa && !llvmcall && nccallargs == 1);
         JL_GC_POP();
-        Value *ptls_pv = emit_bitcast(ctx, ctx.ptlsStates, T_ppjlvalue);
+        Value *ptls_pv = emit_bitcast(ctx, get_current_ptls(ctx), T_ppjlvalue);
         const int nt_offset = offsetof(jl_tls_states_t, next_task);
         Value *pnt = ctx.builder.CreateInBoundsGEP(ptls_pv, ConstantInt::get(T_size, nt_offset / sizeof(void*)));
         ctx.builder.CreateStore(emit_pointer_from_objref(ctx, boxed(ctx, argv[0])), pnt);
@@ -1537,7 +1533,7 @@ static jl_cgval_t emit_ccall(jl_codectx_t &ctx, jl_value_t **args, size_t nargs)
                 checkBB, contBB);
         ctx.builder.SetInsertPoint(checkBB);
         ctx.builder.CreateLoad(
-                ctx.builder.CreateConstInBoundsGEP1_32(T_size, ctx.signalPage, -1),
+                ctx.builder.CreateConstInBoundsGEP1_32(T_size, get_current_signal_page(ctx), -1),
                 true);
         ctx.builder.CreateBr(contBB);
         ctx.f->getBasicBlockList().push_back(contBB);

src/cgutils.cpp

@@ -2761,7 +2761,7 @@ static void emit_cpointercheck(jl_codectx_t &ctx, const jl_cgval_t &x, const std
 // allocation for known size object
 static Value *emit_allocobj(jl_codectx_t &ctx, size_t static_size, Value *jt)
 {
-    Value *ptls_ptr = emit_bitcast(ctx, ctx.ptlsStates, T_pint8);
+    Value *ptls_ptr = emit_bitcast(ctx, get_current_ptls(ctx), T_pint8);
     Function *F = prepare_call(jl_alloc_obj_func);
     auto call = ctx.builder.CreateCall(F, {ptls_ptr, ConstantInt::get(T_size, static_size), maybe_decay_untracked(ctx, jt)});
     call->setAttributes(F->getAttributes());
@@ -3087,7 +3087,7 @@ static void emit_signal_fence(jl_codectx_t &ctx)
 
 static Value *emit_defer_signal(jl_codectx_t &ctx)
 {
-    Value *ptls = emit_bitcast(ctx, ctx.ptlsStates,
+    Value *ptls = emit_bitcast(ctx, get_current_ptls(ctx),
                                         PointerType::get(T_sigatomic, 0));
     Constant *offset = ConstantInt::getSigned(T_int32,
         offsetof(jl_tls_states_t, defer_signal) / sizeof(sig_atomic_t));

src/codegen.cpp

@@ -1084,8 +1084,8 @@ class jl_codectx_t {
     int nargs = 0;
     int nvargs = -1;
 
+    Instruction *currentTask = NULL;
     CallInst *ptlsStates = NULL;
-    Value *signalPage = NULL;
     Value *world_age_field = NULL;
 
     bool debug_enabled = false;
@@ -1121,6 +1121,9 @@ static jl_cgval_t emit_checked_var(jl_codectx_t &ctx, Value *bp, jl_sym_t *name,
 static jl_cgval_t emit_sparam(jl_codectx_t &ctx, size_t i);
 static Value *emit_condition(jl_codectx_t &ctx, const jl_cgval_t &condV, const std::string &msg);
 static void allocate_gc_frame(jl_codectx_t &ctx, BasicBlock *b0);
+static Instruction *get_current_task(jl_codectx_t &ctx);
+static Instruction *get_current_ptls(jl_codectx_t &ctx);
+static Value *get_current_signal_page(jl_codectx_t &ctx);
 static void CreateTrap(IRBuilder<> &irbuilder);
 static CallInst *emit_jlcall(jl_codectx_t &ctx, Function *theFptr, Value *theF,
                              jl_cgval_t *args, size_t nargs, CallingConv::ID cc);
@@ -1194,7 +1197,13 @@ static GlobalVariable *get_pointer_to_constant(jl_codegen_params_t &emission_con
 
 static AllocaInst *emit_static_alloca(jl_codectx_t &ctx, Type *lty)
 {
-    return new AllocaInst(lty, 0, "", /*InsertBefore=*/ctx.ptlsStates);
+    auto InsertBefore = ctx.f->getEntryBlock().getFirstNonPHI();
+    if (InsertBefore) {
+        return new AllocaInst(lty, 0, "", InsertBefore);
+    }
+    else {
+        return new AllocaInst(lty, 0, "", &ctx.f->getEntryBlock());
+    }
 }
 
 static void undef_derived_strct(IRBuilder<> &irbuilder, Value *ptr, jl_datatype_t *sty, MDNode *tbaa)
@@ -4580,17 +4589,66 @@ static void allocate_gc_frame(jl_codectx_t &ctx, BasicBlock *b0)
 
     // allocate a placeholder gc instruction
     ctx.ptlsStates = ctx.builder.CreateCall(prepare_call(jltls_states_func));
-    int nthfield = offsetof(jl_tls_states_t, safepoint) / sizeof(void*);
-    ctx.signalPage = emit_nthptr_recast(ctx, ctx.ptlsStates, nthfield, tbaa_const,
-                                        PointerType::get(T_psize, 0));
+
+    Value *ptls_pv = emit_bitcast(ctx, ctx.ptlsStates, T_pprjlvalue);
+    const int ct_offset = offsetof(jl_tls_states_t, current_task);
+    Value *pct =
+        ctx.builder.CreateInBoundsGEP(ptls_pv,
+                                      ConstantInt::get(T_size, ct_offset / sizeof(void *)),
+                                      "current_task_field");
+    LoadInst *ct =
+        ctx.builder.CreateAlignedLoad(pct, Align(sizeof(void *)), "current_task");
+    tbaa_decorate(tbaa_const, ct);
+    ctx.currentTask = ct;
 }
 
+// Store world age at the entry block of the function. This function should be
+// called right after `allocate_gc_frame` and there should be no context switch.
 static void emit_last_age_field(jl_codectx_t &ctx)
 {
+    auto ptls = ctx.ptlsStates;
+    assert(ctx.builder.GetInsertBlock() == ptls->getParent());
     ctx.world_age_field = ctx.builder.CreateInBoundsGEP(
             T_size,
-            ctx.builder.CreateBitCast(ctx.ptlsStates, T_psize),
-            ConstantInt::get(T_size, offsetof(jl_tls_states_t, world_age) / sizeof(size_t)));
+            ctx.builder.CreateBitCast(ptls, T_psize),
+            ConstantInt::get(T_size, offsetof(jl_tls_states_t, world_age) / sizeof(size_t)),
+            "world_age");
+}
+
+// Get current task.
+static Instruction *get_current_task(jl_codectx_t &ctx)
+{
+    return ctx.currentTask;
+}
+
+// Get PTLS through current task.
+static Instruction *get_current_ptls(jl_codectx_t &ctx)
+{
+    Value *task_pv =
+        emit_bitcast(ctx, decay_derived(ctx, get_current_task(ctx)), T_pprjlvalue);
+    const int ptls_offset = offsetof(jl_task_t, ptls);
+    Value *pptls = ctx.builder.CreateInBoundsGEP(
+        T_prjlvalue, task_pv, ConstantInt::get(T_size, ptls_offset / sizeof(void *)),
+        "ptls_field");
+    LoadInst *ptls_load = ctx.builder.CreateAlignedLoad(
+        emit_bitcast(ctx, pptls, T_ppjlvalue), Align(sizeof(void *)), "ptls_load");
+    // Note: Corersponding store (`t->ptls = ptls`) happes in `ctx_switch` of tasks.c.
+    ptls_load->setOrdering(AtomicOrdering::Monotonic);  // TODO: what should we use?
+    tbaa_decorate(tbaa_const, ptls_load);
+    // Using `CastInst::Create` to get an `Instruction*` without explicit cast:
+    auto ptls = CastInst::Create(Instruction::BitCast, ptls_load, T_ppjlvalue, "ptls");
+    ctx.builder.Insert(ptls);
+    return ptls;
+}
+
+// Get signal page through current task.
+static Value *get_current_signal_page(jl_codectx_t &ctx)
+{
+    // return ctx.builder.CreateCall(prepare_call(reuse_signal_page_func));
+    auto ptls = get_current_ptls(ctx);
+    int nthfield = offsetof(jl_tls_states_t, safepoint) / sizeof(void *);
+    return emit_nthptr_recast(ctx, ptls, nthfield, tbaa_const,
+                              PointerType::get(T_psize, 0));
 }
 
 static Function *emit_tojlinvoke(jl_code_instance_t *codeinst, Module *M, jl_codegen_params_t &params)
@@ -4815,7 +4873,7 @@ static Function* gen_cfun_wrapper(
     emit_last_age_field(ctx);
 
     Value *dummy_world = ctx.builder.CreateAlloca(T_size);
-    Value *have_tls = ctx.builder.CreateIsNotNull(ctx.ptlsStates);
+    Value *have_tls = ctx.builder.CreateIsNotNull(get_current_ptls(ctx));
     // TODO: in the future, try to initialize a full TLS context here
     // for now, just use a dummy field to avoid a branch in this function
     ctx.world_age_field = ctx.builder.CreateSelect(have_tls, ctx.world_age_field, dummy_world);
@@ -6166,9 +6224,9 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
             (va && (int)i == ctx.vaSlot) || // or it's the va arg tuple
             i == 0) { // or it is the first argument (which isn't in `argArray`)
             AllocaInst *av = new AllocaInst(T_prjlvalue, 0,
-                jl_symbol_name(s), /*InsertBefore*/ctx.ptlsStates);
+                jl_symbol_name(s), /*InsertBefore*/get_current_ptls(ctx));
             StoreInst *SI = new StoreInst(V_rnull, av, false, Align(sizeof(void*)));
-            SI->insertAfter(ctx.ptlsStates);
+            SI->insertAfter(get_current_ptls(ctx));
             varinfo.boxroot = av;
             if (ctx.debug_enabled && varinfo.dinfo) {
                 DIExpression *expr;

src/gc.c

@@ -342,8 +342,8 @@ static void jl_gc_push_arraylist(jl_ptls_t ptls, arraylist_t *list)
 {
     void **items = list->items;
     items[0] = (void*)JL_GC_ENCODE_PUSHARGS(list->len - 2);
-    items[1] = ptls->pgcstack;
-    ptls->pgcstack = (jl_gcframe_t*)items;
+    items[1] = jl_pgcstack;
+    jl_pgcstack = (jl_gcframe_t*)items;
 }
 
 // Same assumption as `jl_gc_push_arraylist`. Requires the finalizers lock
@@ -2630,7 +2630,7 @@ mark: {
             uintptr_t lb = 0;
             uintptr_t ub = (uintptr_t)-1;
             if (ptls2 && ta == ptls2->current_task) {
-                s = ptls2->pgcstack;
+                s = ta->gcstack;
             }
             else if (stkbuf) {
                 s = ta->gcstack;

src/init.c

@@ -732,6 +732,15 @@ void _julia_init(JL_IMAGE_SEARCH rel)
 
     jl_gc_enable(0);
 
+    // We need `gcstack` in `Task` to allocate Julia objects; *including* the `Task` type.
+    // However, to allocate a `Task` via `jl_gc_alloc` as done in `jl_init_root_task`,
+    // we need the `Task` type itself. We use stack-allocated "raw" `jl_task_t` struct to
+    // workaround this chicken-and-egg problem. Note that this relies on GC to be turned
+    // off just above as GC fails because we don't/can't allocate the type tag.
+    jl_task_t bootstrap_task;
+    bootstrap_task.gcstack = NULL;
+    jl_current_task = &bootstrap_task;
+
     jl_resolve_sysimg_location(rel);
     // loads sysimg if available, and conditionally sets jl_options.cpu_target
     if (jl_options.image_file)
@@ -749,6 +758,7 @@ void _julia_init(JL_IMAGE_SEARCH rel)
 
     jl_init_tasks();
     jl_init_root_task(stack_lo, stack_hi);
+    assert(jl_current_task != &bootstrap_task);
     jl_init_common_symbols();
     jl_init_flisp();
     jl_init_serializer();

src/julia.h

@@ -735,7 +735,7 @@ typedef struct _jl_gcframe_t {
 // jl_value_t *x=NULL, *y=NULL; JL_GC_PUSH2(&x, &y);
 // x = f(); y = g(); foo(x, y)
 
-#define jl_pgcstack (jl_get_ptls_states()->pgcstack)
+#define jl_pgcstack (jl_current_task->gcstack)
 
 #define JL_GC_ENCODE_PUSHARGS(n)   (((size_t)(n))<<2)
 #define JL_GC_ENCODE_PUSH(n)       ((((size_t)(n))<<2)|1)
@@ -1780,6 +1780,8 @@ typedef struct _jl_task_t {
     unsigned int copy_stack:31; // sizeof stack for copybuf
     unsigned int started:1;
 
+    jl_tls_states_t *ptls;
+
     // saved gc stack top for context switches
     jl_gcframe_t *gcstack;
 } jl_task_t;

src/julia_threads.h

@@ -181,7 +181,6 @@ struct _jl_bt_element_t;
 // Changes to TLS field types must be reflected in codegen.
 #define JL_MAX_BT_SIZE 80000
 struct _jl_tls_states_t {
-    struct _jl_gcframe_t *pgcstack;
     size_t world_age;
     int16_t tid;
     uint64_t rngseed;

src/llvm-final-gc-lowering.cpp

@@ -38,6 +38,7 @@ struct FinalLowerGC: public FunctionPass, private JuliaPassContext {
     Function *poolAllocFunc;
     Function *bigAllocFunc;
     CallInst *ptlsStates;
+    Instruction *pgcstack;
 
     bool doInitialization(Module &M) override;
     bool doFinalization(Module &M) override;
@@ -61,7 +62,7 @@ struct FinalLowerGC: public FunctionPass, private JuliaPassContext {
     // Lowers a `julia.queue_gc_root` intrinsic.
     Value *lowerQueueGCRoot(CallInst *target, Function &F);
 
-    Instruction *getPgcstack(Instruction *ptlsStates);
+    Instruction *getPgcstack();
 };
 
 Value *FinalLowerGC::lowerNewGCFrame(CallInst *target, Function &F)
@@ -111,7 +112,7 @@ void FinalLowerGC::lowerPushGCFrame(CallInst *target, Function &F)
                         T_size->getPointerTo()),
                 Align(sizeof(void*)));
     inst->setMetadata(LLVMContext::MD_tbaa, tbaa_gcframe);
-    Value *pgcstack = builder.Insert(getPgcstack(ptlsStates));
+    Value *pgcstack = getPgcstack();
     inst = builder.CreateAlignedStore(
             builder.CreateAlignedLoad(pgcstack, Align(sizeof(void*))),
             builder.CreatePointerCast(
@@ -139,7 +140,7 @@ void FinalLowerGC::lowerPopGCFrame(CallInst *target, Function &F)
     inst = builder.CreateAlignedStore(
         inst,
         builder.CreateBitCast(
-            builder.Insert(getPgcstack(ptlsStates)),
+            getPgcstack(),
             PointerType::get(T_prjlvalue, 0)),
         Align(sizeof(void*)));
     inst->setMetadata(LLVMContext::MD_tbaa, tbaa_gcframe);
@@ -171,14 +172,28 @@ Value *FinalLowerGC::lowerQueueGCRoot(CallInst *target, Function &F)
     return target;
 }
 
-Instruction *FinalLowerGC::getPgcstack(Instruction *ptlsStates)
+// Get `ptls->current_task->gcstack`
+Instruction *FinalLowerGC::getPgcstack()
 {
-    Constant *offset = ConstantInt::getSigned(T_int32, offsetof(jl_tls_states_t, pgcstack) / sizeof(void*));
-    return GetElementPtrInst::CreateInBounds(
-        T_ppjlvalue,
-        ptlsStates,
-        ArrayRef<Value*>(offset),
-        "jl_pgcstack");
+    if (pgcstack)
+        return pgcstack;
+
+    IRBuilder<> builder(ptlsStates->getParent());
+    if (ptlsStates->getNextNode())
+        builder.SetInsertPoint(ptlsStates->getNextNode());
+    auto task_offset = ConstantInt::getSigned(
+        T_int32, offsetof(jl_tls_states_t, current_task) / sizeof(void *));
+    auto task_field = builder.CreateInBoundsGEP(
+        T_ppjlvalue, ptlsStates, ArrayRef<Value *>(task_offset), "current_task_field");
+    auto task_load =
+        builder.CreateAlignedLoad(task_field, Align(sizeof(void *)), "current_task_load");
+    auto task = builder.CreateBitCast(task_load, T_pppjlvalue, "current_task");
+    auto gcstack_offset =
+        ConstantInt::getSigned(T_int32, offsetof(jl_task_t, gcstack) / sizeof(void *));
+    pgcstack = GetElementPtrInst::CreateInBounds(T_ppjlvalue, task,
+                                                 ArrayRef<Value *>(gcstack_offset));
+    builder.Insert(pgcstack, "jl_pgcstack");
+    return pgcstack;
 }
 
 Value *FinalLowerGC::lowerGCAllocBytes(CallInst *target, Function &F)
@@ -290,6 +305,8 @@ bool FinalLowerGC::runOnFunction(Function &F)
     if (!ptlsStates)
         return true;
 
+    pgcstack = nullptr;
+
     // Acquire intrinsic functions.
     auto newGCFrameFunc = getOrNull(jl_intrinsics::newGCFrame);
     auto pushGCFrameFunc = getOrNull(jl_intrinsics::pushGCFrame);

src/llvm-pass-helpers.cpp

@@ -69,6 +69,7 @@ void JuliaPassContext::initAll(Module &M)
     T_ppjlvalue = PointerType::get(T_pjlvalue, 0);
     T_pjlvalue_der = PointerType::get(T_jlvalue, AddressSpace::Derived);
     T_ppjlvalue_der = PointerType::get(T_prjlvalue, AddressSpace::Derived);
+    T_pppjlvalue = PointerType::get(T_ppjlvalue, 0);
 }
 
 llvm::CallInst *JuliaPassContext::getPtls(llvm::Function &F) const

src/llvm-pass-helpers.h

@@ -49,6 +49,7 @@ struct JuliaPassContext {
     // Types derived from 'jl_value_t'.
     llvm::Type *T_jlvalue;
     llvm::PointerType *T_prjlvalue;
+    llvm::PointerType *T_pppjlvalue;
     llvm::PointerType *T_ppjlvalue;
     llvm::PointerType *T_pjlvalue;
     llvm::PointerType *T_pjlvalue_der;

src/rtutils.c

@@ -214,7 +214,7 @@ JL_DLLEXPORT void jl_enter_handler(jl_handler_t *eh)
     jl_task_t *current_task = ptls->current_task;
     // Must have no safepoint
     eh->prev = current_task->eh;
-    eh->gcstack = ptls->pgcstack;
+    eh->gcstack = current_task->gcstack;
     eh->gc_state = ptls->gc_state;
     eh->locks_len = ptls->locks.len;
     eh->defer_signal = ptls->defer_signal;
@@ -246,7 +246,7 @@ JL_DLLEXPORT void jl_eh_restore_state(jl_handler_t *eh)
     sig_atomic_t old_defer_signal = ptls->defer_signal;
     int8_t old_gc_state = ptls->gc_state;
     current_task->eh = eh->prev;
-    ptls->pgcstack = eh->gcstack;
+    current_task->gcstack = eh->gcstack;
     small_arraylist_t *locks = &ptls->locks;
     int unlocks = locks->len > eh->locks_len;
     if (unlocks) {

src/signals-mach.c

@@ -189,8 +189,9 @@ static void jl_throw_in_thread(int tid, mach_port_t thread, jl_value_t *exceptio
     jl_ptls_t ptls2 = jl_all_tls_states[tid];
     if (!ptls2->safe_restore) {
         assert(exception);
-        ptls2->bt_size = rec_backtrace_ctx(ptls2->bt_data, JL_MAX_BT_SIZE,
-                                           (bt_context_t*)&state, ptls2->pgcstack);
+        ptls2->bt_size =
+            rec_backtrace_ctx(ptls2->bt_data, JL_MAX_BT_SIZE, (bt_context_t *)&state,
+                              ptls2->current_task->gcstack);
         ptls2->sig_exception = exception;
     }
     jl_call_in_state(ptls2, &state, &jl_sig_throw);