optimizer: allow EA-powered finalizer inlining

base/compiler/optimize.jl

@@ -626,47 +626,51 @@ end
 GetNativeEscapeCache(interp::AbstractInterpreter) = GetNativeEscapeCache(code_cache(interp))
 function ((; code_cache)::GetNativeEscapeCache)(mi::MethodInstance)
     codeinst = get(code_cache, mi, nothing)
-    codeinst isa CodeInstance || return false
-    argescapes = traverse_analysis_results(codeinst) do @nospecialize result
+    codeinst isa CodeInstance || return nothing
+    return traverse_analysis_results(codeinst) do @nospecialize result
         return result isa EscapeAnalysis.ArgEscapeCache ? result : nothing
     end
-    if argescapes !== nothing
-        return argescapes
-    end
-    effects = decode_effects(codeinst.ipo_purity_bits)
-    if is_effect_free(effects) && is_inaccessiblememonly(effects)
-        # We might not have run EA on simple frames without any escapes (e.g. when optimization
-        # is skipped when result is constant-folded by abstract interpretation). If those
-        # frames aren't inlined, the accuracy of EA for caller context takes a big hit.
-        # This is a HACK to avoid that, but obviously, a more comprehensive fix would be ideal.
-        return true
-    end
-    return false
+end
+
+analyze_and_cache_escapes!(interp::AbstractInterpreter, opt::OptimizationState, sv::PostOptAnalysisState) =
+    analyze_and_cache_escapes!(interp, opt, sv.ir, sv.result)
+
+function analyze_and_cache_escapes!(interp::AbstractInterpreter, opt::OptimizationState,
+                                    ir::IRCode, result::InferenceResult)
+    nargs = Int(opt.src.nargs)
+    estate = EscapeAnalysis.analyze_escapes(ir, nargs, optimizer_lattice(interp), get_escape_cache(interp))
+    argescapes = EscapeAnalysis.ArgEscapeCache(estate)
+    stack_analysis_result!(result, argescapes)
+    return estate
 end
 
 function refine_effects!(interp::AbstractInterpreter, opt::OptimizationState, sv::PostOptAnalysisState)
+    EA_cached = false
     if !is_effect_free(sv.result.ipo_effects) && sv.all_effect_free && !isempty(sv.ea_analysis_pending)
-        ir = sv.ir
-        nargs = Int(opt.src.nargs)
-        estate = EscapeAnalysis.analyze_escapes(ir, nargs, optimizer_lattice(interp), GetNativeEscapeCache(interp))
-        argescapes = EscapeAnalysis.ArgEscapeCache(estate)
-        stack_analysis_result!(sv.result, argescapes)
+        estate = analyze_and_cache_escapes!(interp, opt, sv)
         validate_mutable_arg_escapes!(estate, sv)
+        EA_cached = true
     end
 
-    any_refinable(sv) || return false
     effects = sv.result.ipo_effects
-    sv.result.ipo_effects = Effects(effects;
+    any_refinable(sv) || @goto run_EA_on_simple_frame
+    effects = sv.result.ipo_effects = Effects(effects;
         consistent = sv.all_retpaths_consistent ? ALWAYS_TRUE : effects.consistent,
         effect_free = sv.all_effect_free ? ALWAYS_TRUE :
-                      sv.effect_free_if_argmem_only === true ? EFFECT_FREE_IF_INACCESSIBLEMEMONLY : effects.effect_free,
+                    sv.effect_free_if_argmem_only === true ? EFFECT_FREE_IF_INACCESSIBLEMEMONLY : effects.effect_free,
         nothrow = sv.all_nothrow ? true : effects.nothrow,
         noub = sv.all_noub ? (sv.any_conditional_ub ? NOUB_IF_NOINBOUNDS : ALWAYS_TRUE) : effects.noub,
         nortcall = sv.nortcall ? true : effects.nortcall)
-    return true
+
+    @label run_EA_on_simple_frame
+    if !EA_cached && is_effect_free(effects) && is_inaccessiblememonly(effects)
+        analyze_and_cache_escapes!(interp, opt, sv)
+    end
+
+    nothing
 end
 
-function is_ipo_dataflow_analysis_profitable(effects::Effects)
+function is_ipo_effects_refinable(effects::Effects)
     return !(is_consistent(effects) && is_effect_free(effects) &&
              is_nothrow(effects) && is_noub(effects))
 end
@@ -941,8 +945,9 @@ end
 
 function ipo_dataflow_analysis!(interp::AbstractInterpreter, opt::OptimizationState,
                                 ir::IRCode, result::InferenceResult)
-    if !is_ipo_dataflow_analysis_profitable(result.ipo_effects)
-        return false
+    if !is_ipo_effects_refinable(result.ipo_effects)
+        analyze_and_cache_escapes!(interp, opt, ir, result)
+        return nothing
     end
 
     @assert isempty(ir.new_nodes) "IRCode should be compacted before post-opt analysis"
@@ -968,7 +973,8 @@ function ipo_dataflow_analysis!(interp::AbstractInterpreter, opt::OptimizationSt
         end
     end
 
-    return refine_effects!(interp, opt, sv)
+    refine_effects!(interp, opt, sv)
+    nothing
 end
 
 # run the optimization work

base/compiler/ssair/EscapeAnalysis/EscapeAnalysis.jl

@@ -18,7 +18,7 @@ import ._TOP_MOD: ==, getindex, setindex!
 using Core: MethodMatch, SimpleVector, ifelse, sizeof
 using Core.IR
 using ._TOP_MOD:     # Base definitions
-    @__MODULE__, @assert, @eval, @goto, @inbounds, @inline, @label, @noinline,
+    @__MODULE__, @assert, @eval, @goto, @inbounds, @inline, @label, @noinline, @show,
     @nospecialize, @specialize, BitSet, Callable, Csize_t, IdDict, IdSet, UnitRange, Vector,
     copy, delete!, empty!, enumerate, error, first, get, get!, haskey, in, isassigned,
     isempty, ismutabletype, keys, last, length, max, min, missing, pop!, push!, pushfirst!,
@@ -657,11 +657,13 @@ function analyze_escapes(ir::IRCode, nargs::Int, 𝕃ₒ::AbstractLattice, get_e
                     # `escape_exception!` conservatively propagates `AllEscape` anyway,
                     # and so escape information imposed on `:the_exception` isn't computed
                     continue
+                elseif head === :gc_preserve_begin
+                    # GC preserve is handled by `escape_gc_preserve!`
+                elseif head === :gc_preserve_end
+                    escape_gc_preserve!(astate, pc, stmt.args)
                 elseif head === :static_parameter ||  # this exists statically, not interested in its escape
-                       head === :copyast ||           # XXX can this account for some escapes?
-                       head === :isdefined ||         # just returns `Bool`, nothing accounts for any escapes
-                       head === :gc_preserve_begin || # `GC.@preserve` expressions themselves won't be used anywhere
-                       head === :gc_preserve_end      # `GC.@preserve` expressions themselves won't be used anywhere
+                       head === :copyast ||           # XXX escape something?
+                       head === :isdefined            # just returns `Bool`, nothing accounts for any escapes
                     continue
                 else
                     add_conservative_changes!(astate, pc, stmt.args)
@@ -1064,16 +1066,10 @@ end
 function escape_invoke!(astate::AnalysisState, pc::Int, args::Vector{Any})
     mi = first(args)::MethodInstance
     first_idx, last_idx = 2, length(args)
+    add_liveness_changes!(astate, pc, args, first_idx, last_idx)
     # TODO inspect `astate.ir.stmts[pc][:info]` and use const-prop'ed `InferenceResult` if available
     cache = astate.get_escape_cache(mi)
-    if cache isa Bool
-        if cache
-            return nothing # guaranteed to have no escape
-        else
-            return add_conservative_changes!(astate, pc, args, 2)
-        end
-    end
-    cache = cache::ArgEscapeCache
+    cache isa ArgEscapeCache || return add_conservative_changes!(astate, pc, args, 2)
     ret = SSAValue(pc)
     retinfo = astate.estate[ret] # escape information imposed on the call statement
     method = mi.def::Method
@@ -1162,6 +1158,17 @@ function escape_foreigncall!(astate::AnalysisState, pc::Int, args::Vector{Any})
     end
 end
 
+function escape_gc_preserve!(astate::AnalysisState, pc::Int, args::Vector{Any})
+    @assert length(args) == 1 "invalid :gc_preserve_end"
+    val = args[1]
+    @assert val isa SSAValue "invalid :gc_preserve_end"
+    beginstmt = astate.ir[val][:stmt]
+    @assert isexpr(beginstmt, :gc_preserve_begin) "invalid :gc_preserve_end"
+    beginargs = beginstmt.args
+    # COMBAK we might need to add liveness for all statements from `:gc_preserve_begin` to `:gc_preserve_end`
+    add_liveness_changes!(astate, pc, beginargs)
+end
+
 normalize(@nospecialize x) = isa(x, QuoteNode) ? x.value : x
 
 function escape_call!(astate::AnalysisState, pc::Int, args::Vector{Any})
@@ -1187,20 +1194,12 @@ function escape_call!(astate::AnalysisState, pc::Int, args::Vector{Any})
     if result === missing
         # if this call hasn't been handled by any of pre-defined handlers, escape it conservatively
         add_conservative_changes!(astate, pc, args)
-        return
     elseif result === true
         add_liveness_changes!(astate, pc, args, 2)
-        return # ThrownEscape is already checked
+    elseif is_nothrow(astate.ir, pc)
+        add_liveness_changes!(astate, pc, args, 2)
     else
-        # we escape statements with the `ThrownEscape` property using the effect-freeness
-        # computed by `stmt_effect_flags` invoked within inlining
-        # TODO throwness ≠ "effect-free-ness"
-        if is_nothrow(astate.ir, pc)
-            add_liveness_changes!(astate, pc, args, 2)
-        else
-            add_fallback_changes!(astate, pc, args, 2)
-        end
-        return
+        add_fallback_changes!(astate, pc, args, 2)
     end
 end
 
@@ -1528,4 +1527,12 @@ function escape_array_copy!(astate::AnalysisState, pc::Int, args::Vector{Any})
     add_liveness_changes!(astate, pc, args, 6)
 end
 
+function escape_builtin!(::typeof(Core.finalizer), astate::AnalysisState, pc::Int, args::Vector{Any})
+    if length(args) ≥ 3
+        obj = args[3]
+        add_liveness_change!(astate, obj, pc) # TODO setup a proper FinalizerEscape?
+    end
+    return false
+end
+
 end # baremodule EscapeAnalysis

base/compiler/ssair/passes.jl

@@ -1300,7 +1300,13 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
                 # Inlining performs legality checks on the finalizer to determine
                 # whether or not we may inline it. If so, it appends extra arguments
                 # at the end of the intrinsic. Detect that here.
-                length(stmt.args) == 5 || continue
+                if length(stmt.args) == 4 && stmt.args[4] === nothing
+                    # constant case
+                elseif length(stmt.args) == 5 && stmt.args[4] isa Bool && stmt.args[5] isa MethodInstance
+                    # inlining case
+                else
+                    continue
+                end
             end
             is_finalizer = true
         elseif isexpr(stmt, :foreigncall)
@@ -1685,18 +1691,21 @@ end
 function sroa_mutables!(ir::IRCode, defuses::IdDict{Int,Tuple{SPCSet,SSADefUse}}, used_ssas::Vector{Int}, lazydomtree::LazyDomtree, inlining::Union{Nothing,InliningState})
     𝕃ₒ = inlining === nothing ? SimpleInferenceLattice.instance : optimizer_lattice(inlining.interp)
     lazypostdomtree = LazyPostDomtree(ir)
-    for (defidx, (intermediaries, defuse)) in defuses
-        # Check if there are any uses we did not account for. If so, the variable
-        # escapes and we cannot eliminate the allocation. This works, because we're guaranteed
-        # not to include any intermediaries that have dead uses. As a result, missing uses will only ever
-        # show up in the nuses_total count.
-        nleaves = length(defuse.uses) + length(defuse.defs)
-        nuses = 0
-        for iidx in intermediaries
-            nuses += used_ssas[iidx]
+    function find_finalizer_useidx(defuse::SSADefUse)
+        finalizer_useidx = nothing
+        for (useidx, use) in enumerate(defuse.uses)
+            if use.kind === :finalizer
+                # For now: Only allow one finalizer per allocation
+                finalizer_useidx !== nothing && return false
+                finalizer_useidx = useidx
+            end
         end
-        nuses_total = used_ssas[defidx] + nuses - length(intermediaries)
-        nleaves == nuses_total || continue
+        if finalizer_useidx === nothing || inlining === nothing
+            return true
+        end
+        return finalizer_useidx
+    end
+    for (defidx, (intermediaries, defuse)) in defuses
         # Find the type for this allocation
         defexpr = ir[SSAValue(defidx)][:stmt]
         isexpr(defexpr, :new) || continue
@@ -1706,22 +1715,47 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int,Tuple{SPCSet,SSADefUse}}
         typ = widenconst(typ)
         ismutabletype(typ) || continue
         typ = typ::DataType
-        # First check for any finalizer calls
-        finalizer_useidx = nothing
-        for (useidx, use) in enumerate(defuse.uses)
-            if use.kind === :finalizer
-                # For now: Only allow one finalizer per allocation
-                finalizer_useidx !== nothing && @goto skip
-                finalizer_useidx = useidx
-            end
+        # Check if there are any uses we did not account for. If so, the variable
+        # escapes and we cannot eliminate the allocation. This works, because we're guaranteed
+        # not to include any intermediaries that have dead uses. As a result, missing uses will only ever
+        # show up in the nuses_total count.
+        nleaves = length(defuse.uses) + length(defuse.defs)
+        nuses = 0
+        for iidx in intermediaries
+            nuses += used_ssas[iidx]
         end
+        nuses_total = used_ssas[defidx] + nuses - length(intermediaries)
         all_eliminated = all_forwarded = true
-        if finalizer_useidx !== nothing && inlining !== nothing
-            finalizer_idx = defuse.uses[finalizer_useidx].idx
-            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
+        if nleaves ≠ nuses_total
+            finalizer_useidx = find_finalizer_useidx(defuse)
+            if finalizer_useidx isa Int
+                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)
+                    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,
+                        lazydomtree, lazypostdomtree, ir[SSAValue(finalizer_idx)][:info])
+                end
+            end
+            continue
+        else
+            finalizer_useidx = find_finalizer_useidx(defuse)
+            if finalizer_useidx isa Int
+                finalizer_idx = defuse.uses[finalizer_useidx].idx
+                try_resolve_finalizer!(ir, defidx, finalizer_idx, defuse, inlining::InliningState,
+                    lazydomtree, lazypostdomtree, ir[SSAValue(finalizer_idx)][:info])
+                deleteat!(defuse.uses, finalizer_useidx)
+                all_eliminated = all_forwarded = false # can't eliminate `setfield!` calls safely
+            elseif !finalizer_useidx
+                continue
+            end
         end
         # Partition defuses by field
         fielddefuse = SSADefUse[SSADefUse() for _ = 1:fieldcount(typ)]

base/compiler/types.jl

@@ -455,6 +455,8 @@ typeinf_lattice(::AbstractInterpreter) = InferenceLattice(BaseInferenceLattice.i
 ipo_lattice(::AbstractInterpreter) = InferenceLattice(IPOResultLattice.instance)
 optimizer_lattice(::AbstractInterpreter) = SimpleInferenceLattice.instance
 
+get_escape_cache(interp::AbstractInterpreter) = GetNativeEscapeCache(interp)
+
 abstract type CallInfo end
 
 @nospecialize