optimizer: Add early finalize calls

base/compiler/optimize.jl

@@ -87,7 +87,7 @@ function stmt_effect_free end # imported by EscapeAnalysis
 function alloc_array_ndims end # imported by EscapeAnalysis
 include("compiler/ssair/driver.jl")
 using .EscapeAnalysis
-import .EscapeAnalysis: EscapeState, ArgEscapeCache, is_ipo_profitable
+import .EscapeAnalysis: EscapeState, EscapeInfo, ArgEscapeCache, is_ipo_profitable
 
 """
     cache_escapes!(caller::InferenceResult, estate::EscapeState)
@@ -112,6 +112,73 @@ function ipo_escape_cache(mi_cache::MICache) where MICache
 end
 null_escape_cache(linfo::Union{InferenceResult,MethodInstance}) = nothing
 
+"""
+    early_finalize!(ir::IRCode, estate::EscapeState, domtree::DomTree) -> IRCode
+
+Analyzes `ir` for heap allocations which escape only via `FinalizerEscape`
+(thus having a `finalizer` call associated), and inserts a call to
+`finalize(obj)` just before any return where the allocation doesn't escape.
+"""
+function early_finalize!(ir::IRCode, estate::EscapeState, domtree::DomTree)
+    # Find all allocations that escape only through a finalizer
+    to_finalize = Pair{Int,EscapeInfo}[]
+    for idx in (estate.nargs+1):length(estate.escapes)
+        info = estate.escapes[idx]
+        if has_finalizer_escape(info)# && !has_return_escape(info)
+            push!(to_finalize, (idx-estate.nargs)=>info)
+        end
+    end
+
+    # Find all locations to insert calls to `finalize`
+    # currently this is just all returns
+    return_locs = Int[]
+    for bb in ir.cfg.blocks
+        inst = ir.stmts.inst[bb.stmts.stop]
+        if inst isa ReturnNode
+            push!(return_locs, bb.stmts.stop)
+        end
+    end
+
+    # Insert `finalize` calls in order
+    isalloc(ir::IRCode, pc::Int) = isexpr(ir.stmts[pc][:inst], :new)
+    for idx in return_locs
+        non_escaping = Pair{Int,EscapeInfo}[]
+        for (alloc_idx, alloc_info) in to_finalize
+            allblocks = Int[]
+            fdu = FieldDefUse()
+            for pc in alloc_info.Liveness
+                pc < 1 && continue
+                real_pc = pc - estate.nargs
+                @assert 1 <= real_pc <= length(ir.stmts)
+                if isalloc(ir, real_pc)
+                    push!(fdu.defs, real_pc)
+                    push!(allblocks, block_for_inst(ir, real_pc))
+                end
+            end
+            # Not necessary: `push!(fdu.uses, idx)`
+            use = idx
+            defs = find_def_for_use(ir, domtree, allblocks, fdu, use)
+            if !has_return_escape(alloc_info, idx) &&
+               !has_arg_escape(alloc_info) &&
+               defs !== nothing
+                push!(non_escaping, alloc_idx=>alloc_info)
+            end
+        end
+        if length(non_escaping) > 0
+            ct_expr = Expr(:foreigncall, :(:jl_get_current_task), Ref{Task}, svec(), 0, :(:ccall))
+            ct_insn = NewInstruction(ct_expr, Ref{Task})
+            ct = insert_node!(ir, idx, ct_insn)
+            for (alloc_idx, alloc_info) in non_escaping
+                fin_expr = Expr(:foreigncall, :(:jl_finalize_th), Nothing, svec(Any, Any), 0, :(:ccall), ct, SSAValue(alloc_idx))
+                fin_insn = NewInstruction(fin_expr, Nothing)
+                fin = insert_node!(ir, idx, fin_insn)
+            end
+        end
+    end
+
+    ir
+end
+
 mutable struct OptimizationState
     linfo::MethodInstance
     src::CodeInfo

base/compiler/ssair/passes.jl

@@ -582,7 +582,13 @@ function linear_pass!(ir::IRCode)
                 elseif isexpr(def, :new)
                     typ = unwrap_unionall(widenconst(argextype(SSAValue(defidx), compact)))
                     if typ isa DataType
-                        ismutabletype(typ) && continue # mutable SROA is performed later
+                        if ismutabletype(typ)
+                            if stmt.args[1] isa QuoteNode && stmt.args[1].value == :jl_gc_add_finalizer_th
+                                # Track this for `early_finalize!`
+                                anymutability = true
+                            end
+                            continue # mutable SROA is performed later
+                        end
                         record_immutable_preserve!(new_preserves, def, compact)
                         push!(preserved, preserved_arg.id)
                     end
@@ -883,6 +889,8 @@ function memory_opt_pass!(ir::IRCode, estate::EscapeState)
         end
     end
 
+    @timeit "Early Finalize" ir = early_finalize!(ir, estate, domtree)
+
     return ir
 end