optimizer: support early finalization

base/Base.jl

@@ -277,8 +277,7 @@ end
 
 BUILDROOT::String = ""
 
-baremodule BuildSettings
-end
+baremodule BuildSettings end
 
 let i = 1
     global BUILDROOT

base/compiler/abstractinterpretation.jl

@@ -2280,10 +2280,10 @@ function abstract_finalizer(interp::AbstractInterpreter, argtypes::Vector{Any},
         finalizer_argvec = Any[argtypes[2], argtypes[3]]
         call = abstract_call(interp, ArgInfo(nothing, finalizer_argvec), StmtInfo(false), sv, #=max_methods=#1)::Future
         return Future{CallMeta}(call, interp, sv) do call, interp, sv
-            return CallMeta(Nothing, Any, Effects(), FinalizerInfo(call.info, call.effects))
+            return CallMeta(Int, Any, Effects(), FinalizerInfo(call.info, call.effects))
         end
     end
-    return Future(CallMeta(Nothing, Any, Effects(), NoCallInfo()))
+    return Future(CallMeta(Int, Any, Effects(), NoCallInfo()))
 end
 
 function abstract_throw(interp::AbstractInterpreter, argtypes::Vector{Any}, ::AbsIntState)

base/compiler/ssair/passes.jl

@@ -1594,11 +1594,11 @@ function try_resolve_finalizer!(ir::IRCode, alloc_idx::Int, finalizer_idx::Int,
         inlining::InliningState, lazydomtree::LazyDomtree,
         lazypostdomtree::LazyPostDomtree, @nospecialize(info::CallInfo))
     # For now, require that:
-    # 1. The allocation dominates the finalizer registration
+    # 1. The allocation dominates the finalizer registration.
     # 2. The insertion block for the finalizer is the post-dominator of all
     #    uses (including the finalizer registration).
     # 3. The path from the finalizer registration to the finalizer inlining
-    #    location is nothrow
+    #    location is nothrow, otherwise we need to insert `Core._cancel_finalizer` too.
     #
     # TODO: We could relax the check 2, by inlining the finalizer multiple times.
 
@@ -1631,19 +1631,16 @@ function try_resolve_finalizer!(ir::IRCode, alloc_idx::Int, finalizer_idx::Int,
     foreach(note_defuse!, defuse.defs)
     insert_bb != 0 || return nothing # verify post-dominator of all uses exists
 
+    # Figure out the exact statement where we're going to inline the finalizer.
+    finalizer_stmt = ir[SSAValue(finalizer_idx)][:stmt]
+
+    current_task_ssa = nothing
     if !OptimizationParams(inlining.interp).assume_fatal_throw
         # Collect all reachable blocks between the finalizer registration and the
         # insertion point
         blocks = reachable_blocks(ir.cfg, finalizer_bb, insert_bb)
 
         # Check #3
-        function check_range_nothrow(s::Int, e::Int)
-            return all(s:e) do sidx::Int
-                sidx == finalizer_idx && return true
-                sidx == alloc_idx && return true
-                return is_nothrow(ir, SSAValue(sidx))
-            end
-        end
         for bb in blocks
             range = ir.cfg.blocks[bb].stmts
             s, e = first(range), last(range)
@@ -1654,18 +1651,40 @@ function try_resolve_finalizer!(ir::IRCode, alloc_idx::Int, finalizer_idx::Int,
             if bb == finalizer_bb
                 s = finalizer_idx
             end
-            check_range_nothrow(s, e) || return nothing
+            all(s:e) do sidx::Int
+                sidx == finalizer_idx && return true
+                sidx == alloc_idx && return true
+                return is_nothrow(ir, SSAValue(sidx))
+            end && continue
+
+            # An exception may be thrown between the finalizer registration and the point
+            # where the object’s lifetime ends (`insert_idx`): In such cases, we can’t
+            # remove the finalizer registration, but we can still inline the finalizer
+            # with inserting `Core._cancel_finalizer` at the end.
+            # Here, prepare a reference to the current task object that should be passed to
+            # `Core._cancel_finalizer` and insert it into `Core.finalizer` so that the
+            # finalizer is added to the ptls of the current task.
+            current_task_stmt = Expr(:foreigncall, QuoteNode(:jl_get_current_task),
+                Core.Ref{Core.Task}, Core.svec(), 0, QuoteNode(:ccall))
+            newinst = NewInstruction(current_task_stmt, Core.Task)
+            current_task_ssa = insert_node!(ir, finalizer_idx, newinst)
+            push!(finalizer_stmt.args, current_task_ssa)
+            break
         end
     end
 
-    # Ok, legality check complete. Figure out the exact statement where we're
-    # going to inline the finalizer.
     loc = insert_idx === nothing ? first(ir.cfg.blocks[insert_bb].stmts) : insert_idx::Int
     attach_after = insert_idx !== nothing
-
-    finalizer_stmt = ir[SSAValue(finalizer_idx)][:stmt]
-    argexprs = Any[finalizer_stmt.args[2], finalizer_stmt.args[3]]
     flag = info isa FinalizerInfo ? flags_for_effects(info.effects) : IR_FLAG_NULL
+    alloc_obj = finalizer_stmt.args[3]
+    cancellation_required = current_task_ssa !== nothing
+    if cancellation_required
+        lookup_idx_ssa = SSAValue(finalizer_idx)
+        finalize_call = Expr(:call, GlobalRef(Core, :_cancel_finalizer), alloc_obj, current_task_ssa, lookup_idx_ssa)
+        newinst = add_flag(NewInstruction(finalize_call, Nothing), flag)
+        insert_node!(ir, loc, newinst, attach_after)
+    end
+    argexprs = Any[finalizer_stmt.args[2], alloc_obj]
     if length(finalizer_stmt.args) >= 4
         inline = finalizer_stmt.args[4]
         if inline === nothing
@@ -1683,8 +1702,10 @@ function try_resolve_finalizer!(ir::IRCode, alloc_idx::Int, finalizer_idx::Int,
         newinst = add_flag(NewInstruction(Expr(:call, argexprs...), Nothing), flag)
         insert_node!(ir, loc, newinst, attach_after)
     end
-    # Erase the call to `finalizer`
-    ir[SSAValue(finalizer_idx)][:stmt] = nothing
+    if !cancellation_required
+        # Erase the call to `finalizer`
+        ir[SSAValue(finalizer_idx)][:stmt] = nothing
+    end
     return nothing
 end
 
@@ -1700,10 +1721,7 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int,Tuple{SPCSet,SSADefUse}}
                 finalizer_useidx = useidx
             end
         end
-        if finalizer_useidx === nothing || inlining === nothing
-            return true
-        end
-        return finalizer_useidx
+        return something(finalizer_useidx, true)
     end
     for (defidx, (intermediaries, defuse)) in defuses
         # Find the type for this allocation
@@ -1728,28 +1746,28 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int,Tuple{SPCSet,SSADefUse}}
         all_eliminated = all_forwarded = true
         if nleaves ≠ nuses_total
             finalizer_useidx = find_finalizer_useidx(defuse)
-            if finalizer_useidx isa Int
+            if finalizer_useidx isa Int && inlining !== nothing
                 nargs = length(ir.argtypes) # COMBAK this might need to be `Int(opt.src.nargs)`
                 estate = EscapeAnalysis.analyze_escapes(ir, nargs, 𝕃ₒ, get_escape_cache(inlining.interp))
                 einfo = estate[SSAValue(defidx)]
-                if EscapeAnalysis.has_no_escape(einfo)
+                if EscapeAnalysis.has_no_escape(EscapeAnalysis.ignore_thrownescapes(einfo))
                     already = BitSet(use.idx for use in defuse.uses)
                     for idx = einfo.Liveness
                         if idx ∉ already
                             push!(defuse.uses, SSAUse(:EALiveness, idx))
                         end
                     end
                     finalizer_idx = defuse.uses[finalizer_useidx].idx
-                    try_resolve_finalizer!(ir, defidx, finalizer_idx, defuse, inlining::InliningState,
+                    try_resolve_finalizer!(ir, defidx, finalizer_idx, defuse, inlining,
                         lazydomtree, lazypostdomtree, ir[SSAValue(finalizer_idx)][:info])
                 end
             end
             continue
         else
             finalizer_useidx = find_finalizer_useidx(defuse)
-            if finalizer_useidx isa Int
+            if finalizer_useidx isa Int && inlining !== nothing
                 finalizer_idx = defuse.uses[finalizer_useidx].idx
-                try_resolve_finalizer!(ir, defidx, finalizer_idx, defuse, inlining::InliningState,
+                try_resolve_finalizer!(ir, defidx, finalizer_idx, defuse, inlining,
                     lazydomtree, lazypostdomtree, ir[SSAValue(finalizer_idx)][:info])
                 deleteat!(defuse.uses, finalizer_useidx)
                 all_eliminated = all_forwarded = false # can't eliminate `setfield!` calls safely

base/compiler/tfuncs.jl

@@ -750,7 +750,13 @@ add_tfunc(donotdelete, 0, INT_INF, @nospecs((𝕃::AbstractLattice, args...)->No
     end
 end
 add_tfunc(compilerbarrier, 2, 2, compilerbarrier_tfunc, 5)
-add_tfunc(Core.finalizer, 2, 4, @nospecs((𝕃::AbstractLattice, args...)->Nothing), 5)
+add_tfunc(Core.finalizer, 2, 5, @nospecs((𝕃::AbstractLattice, args...)->Int), 5)
+@nospecs function _cancel_finalizer_tfunc(𝕃::AbstractLattice, o, ct, lookup_idx)
+    hasintersect(widenconst(ct), Task) || return Bottom
+    hasintersect(widenconst(lookup_idx), Int) || return Bottom
+    return Nothing
+end
+add_tfunc(Core._cancel_finalizer, 3, 3, _cancel_finalizer_tfunc, 5)
 
 @nospecs function compilerbarrier_nothrow(setting, val)
     return isa(setting, Const) && contains_is((:type, :const, :conditional), setting.val)
@@ -2297,9 +2303,12 @@ function _builtin_nothrow(𝕃::AbstractLattice, @nospecialize(f::Builtin), argt
     elseif f === donotdelete
         return true
     elseif f === Core.finalizer
-        2 <= na <= 4 || return false
-        # Core.finalizer does no error checking - that's done in Base.finalizer
+        2 <= na <= 5 || return false
+        # `Core.finalizer` does no error checking - that's done in Base.finalizer
         return true
+    elseif f === Core._cancel_finalizer
+        na == 3 || return false
+        return argtypes[2] ⊑ Task && argtypes[3] ⊑ Int
     elseif f === Core.compilerbarrier
         na == 2 || return false
         return compilerbarrier_nothrow(argtypes[1], nothing)

base/gcutils.jl

@@ -1,6 +1,5 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
-
 """
     WeakRef(x)
 
@@ -99,8 +98,7 @@ end
 
 Immediately run finalizers registered for object `x`.
 """
-finalize(@nospecialize(o)) = ccall(:jl_finalize_th, Cvoid, (Any, Any,),
-                                   current_task(), o)
+finalize(@nospecialize(o)) = ccall(:jl_finalize_th, Cvoid, (Any, Any,), current_task(), o)
 
 """
     Base.GC

src/builtin_proto.h

@@ -26,6 +26,7 @@ DECLARE_BUILTIN(_apply_pure);
 DECLARE_BUILTIN(_call_in_world);
 DECLARE_BUILTIN(_call_in_world_total);
 DECLARE_BUILTIN(_call_latest);
+DECLARE_BUILTIN(_cancel_finalizer);
 DECLARE_BUILTIN(_compute_sparams);
 DECLARE_BUILTIN(_expr);
 DECLARE_BUILTIN(_svec_ref);

src/builtins.c

@@ -2004,9 +2004,25 @@ JL_CALLABLE(jl_f_compilerbarrier)
 JL_CALLABLE(jl_f_finalizer)
 {
     // NOTE the compiler may temporarily insert additional argument for the later inlining pass
-    JL_NARGS(finalizer, 2, 4);
+    JL_NARGS(finalizer, 2, 5);
+    if (nargs == 5 && jl_is_task(args[4]))
+        // There are cases where the compiler inserts `current_task` as the 5th argument,
+        // and in such cases, the finalizer is added to the `ptls` of that task.
+        // This task is later referenced by `_cancel_finalizer`.
+        return jl_box_long(jl_gc_add_finalizer_(((jl_task_t*)args[4])->ptls, args[1], args[0]));
     jl_task_t *ct = jl_current_task;
-    jl_gc_add_finalizer_(ct->ptls, args[1], args[0]);
+    return jl_box_long(jl_gc_add_finalizer_(ct->ptls, args[1], args[0]));
+}
+
+JL_CALLABLE(jl_f__cancel_finalizer)
+{
+    JL_NARGS(_cancel_finalizer, 3, 3);
+    JL_TYPECHK(_cancel_finalizer, task, args[1])
+    JL_TYPECHK(_cancel_finalizer, long, args[2])
+    jl_value_t *o = args[0];
+    jl_task_t *ct = (jl_task_t*)args[1];
+    size_t lookup_idx = jl_unbox_long(args[2]);
+    jl_cancel_finalizer(o, ct, lookup_idx);
     return jl_nothing;
 }
 
@@ -2442,6 +2458,7 @@ void jl_init_primitives(void) JL_GC_DISABLED
     jl_builtin_donotdelete = add_builtin_func("donotdelete", jl_f_donotdelete);
     jl_builtin_compilerbarrier = add_builtin_func("compilerbarrier", jl_f_compilerbarrier);
     add_builtin_func("finalizer", jl_f_finalizer);
+    add_builtin_func("_cancel_finalizer", jl_f__cancel_finalizer);
     add_builtin_func("_compute_sparams", jl_f__compute_sparams);
     add_builtin_func("_svec_ref", jl_f__svec_ref);
     jl_builtin_current_scope = add_builtin_func("current_scope", jl_f_current_scope);