Sink CodeInfo transformation into transform_result_for_cache, continued (JuliaLang#57375)

Continuation of the work started in JuliaLang#56897.

Co-authored-by: Cédric Belmant <[email protected]>

Author: serenity4

extras/CompilerDevTools/src/CompilerDevTools.jl

@@ -47,7 +47,7 @@ end
 
 function Compiler.optimize(interp::SplitCacheInterp, opt::Compiler.OptimizationState, caller::Compiler.InferenceResult)
     @invoke Compiler.optimize(interp::Compiler.AbstractInterpreter, opt::Compiler.OptimizationState, caller::Compiler.InferenceResult)
-    ir = opt.ir::Compiler.IRCode
+    ir = opt.result.ir::Compiler.IRCode
     override = GlobalRef(@__MODULE__(), :with_new_compiler)
     for inst in ir.stmts
         stmt = inst[:stmt]

src/optimize.jl

@@ -129,6 +129,7 @@ is_declared_noinline(@nospecialize src::MaybeCompressed) =
 # return whether this src should be inlined. If so, retrieve_ir_for_inlining must return an IRCode from it
 function src_inlining_policy(interp::AbstractInterpreter,
     @nospecialize(src), @nospecialize(info::CallInfo), stmt_flag::UInt32)
+    isa(src, OptimizationState) && (src = src.src)
     if isa(src, MaybeCompressed)
         src_inlineable = is_stmt_inline(stmt_flag) || is_inlineable(src)
         return src_inlineable
@@ -154,10 +155,20 @@ end
 # get `code_cache(::AbstractInterpreter)` from `state::InliningState`
 code_cache(state::InliningState) = WorldView(code_cache(state.interp), state.world)
 
+mutable struct OptimizationResult
+    ir::IRCode
+    simplified::Bool # indicates whether the IR was processed with `cfg_simplify!`
+end
+
+function simplify_ir!(result::OptimizationResult)
+    result.ir = cfg_simplify!(result.ir)
+    result.simplified = true
+end
+
 mutable struct OptimizationState{Interp<:AbstractInterpreter}
     linfo::MethodInstance
     src::CodeInfo
-    ir::Union{Nothing, IRCode}
+    result::Union{Nothing, OptimizationResult}
     stmt_info::Vector{CallInfo}
     mod::Module
     sptypes::Vector{VarState}
@@ -226,10 +237,13 @@ include("ssair/passes.jl")
 include("ssair/irinterp.jl")
 
 function ir_to_codeinf!(opt::OptimizationState)
-    (; linfo, src) = opt
-    src = ir_to_codeinf!(src, opt.ir::IRCode)
-    src.edges = Core.svec(opt.inlining.edges...)
-    opt.ir = nothing
+    (; linfo, src, result) = opt
+    if result === nothing
+        return src
+    end
+    src = ir_to_codeinf!(src, result.ir)
+    opt.result = nothing
+    opt.src = src
     maybe_validate_code(linfo, src, "optimized")
     return src
 end
@@ -489,7 +503,7 @@ function finish(interp::AbstractInterpreter, opt::OptimizationState,
     @assert !(result isa LimitedAccuracy)
     result = widenslotwrapper(result)
 
-    opt.ir = ir
+    opt.result = OptimizationResult(ir, false)
 
     # determine and cache inlineability
     if !force_noinline

src/ssair/inlining.jl

@@ -975,6 +975,14 @@ function retrieve_ir_for_inlining(mi::MethodInstance, ir::IRCode, preserve_local
     ir.debuginfo.def = mi
     return ir, spec_info, DebugInfo(ir.debuginfo, length(ir.stmts))
 end
+function retrieve_ir_for_inlining(mi::MethodInstance, opt::OptimizationState, preserve_local_sources::Bool)
+    result = opt.result
+    if result !== nothing
+        !result.simplified && simplify_ir!(result)
+        return retrieve_ir_for_inlining(mi, result.ir, preserve_local_sources)
+    end
+    retrieve_ir_for_inlining(mi, opt.src, preserve_local_sources)
+end
 
 function handle_single_case!(todo::Vector{Pair{Int,Any}},
     ir::IRCode, idx::Int, stmt::Expr, @nospecialize(case),

src/ssair/ir.jl

@@ -1265,39 +1265,48 @@ function process_phinode_values(old_values::Vector{Any}, late_fixup::Vector{Int}
     values = Vector{Any}(undef, length(old_values))
     for i = 1:length(old_values)
         isassigned(old_values, i) || continue
-        val = old_values[i]
-        if isa(val, SSAValue)
-            if do_rename_ssa
-                if !already_inserted(i, OldSSAValue(val.id))
-                    push!(late_fixup, result_idx)
-                    val = OldSSAValue(val.id)
-                else
-                    val = renumber_ssa2(val, ssa_rename, used_ssas, new_new_used_ssas, do_rename_ssa, mark_refined!)
-                end
-            else
-                used_ssas[val.id] += 1
-            end
-        elseif isa(val, OldSSAValue)
-            if !already_inserted(i, val)
+        values[i] = process_phinode_value(old_values, i, late_fixup, already_inserted, result_idx, ssa_rename, used_ssas, new_new_used_ssas, do_rename_ssa, mark_refined!)
+    end
+    return values
+end
+
+function process_phinode_value(old_values::Vector{Any}, i::Int, late_fixup::Vector{Int},
+                               already_inserted, result_idx::Int,
+                               ssa_rename::Vector{Any}, used_ssas::Vector{Int},
+                               new_new_used_ssas::Vector{Int},
+                               do_rename_ssa::Bool,
+                               mark_refined!::Union{Refiner, Nothing})
+    val = old_values[i]
+    if isa(val, SSAValue)
+        if do_rename_ssa
+            if !already_inserted(i, OldSSAValue(val.id))
                 push!(late_fixup, result_idx)
+                val = OldSSAValue(val.id)
             else
-                # Always renumber these. do_rename_ssa applies only to actual SSAValues
-                val = renumber_ssa2(SSAValue(val.id), ssa_rename, used_ssas, new_new_used_ssas, true, mark_refined!)
-            end
-        elseif isa(val, NewSSAValue)
-            if val.id < 0
-                new_new_used_ssas[-val.id] += 1
-            else
-                @assert do_rename_ssa
-                val = SSAValue(val.id)
+                val = renumber_ssa2(val, ssa_rename, used_ssas, new_new_used_ssas, do_rename_ssa, mark_refined!)
             end
+        else
+            used_ssas[val.id] += 1
         end
-        if isa(val, NewSSAValue)
+    elseif isa(val, OldSSAValue)
+        if !already_inserted(i, val)
             push!(late_fixup, result_idx)
+        else
+            # Always renumber these. do_rename_ssa applies only to actual SSAValues
+            val = renumber_ssa2(SSAValue(val.id), ssa_rename, used_ssas, new_new_used_ssas, true, mark_refined!)
+        end
+    elseif isa(val, NewSSAValue)
+        if val.id < 0
+            new_new_used_ssas[-val.id] += 1
+        else
+            @assert do_rename_ssa
+            val = SSAValue(val.id)
         end
-        values[i] = val
     end
-    return values
+    if isa(val, NewSSAValue)
+        push!(late_fixup, result_idx)
+    end
+    return val
 end
 
 function renumber_ssa2(val::SSAValue, ssanums::Vector{Any}, used_ssas::Vector{Int},

src/ssair/passes.jl

@@ -2574,51 +2574,50 @@ function cfg_simplify!(ir::IRCode)
                     values = phi.values
                     (; ssa_rename, late_fixup, used_ssas, new_new_used_ssas) = compact
                     ssa_rename[i] = SSAValue(compact.result_idx)
-                    already_inserted = function (i::Int, val::OldSSAValue)
+                    already_inserted = function (branch::Int, val::OldSSAValue)
                         if val.id in old_bb_stmts
                             return val.id <= i
                         end
-                        return bb_rename_pred[phi.edges[i]] < idx
+                        return 0 < bb_rename_pred[phi.edges[branch]] < idx
                     end
-                    renamed_values = process_phinode_values(values, late_fixup, already_inserted, compact.result_idx, ssa_rename, used_ssas, new_new_used_ssas, true, nothing)
                     edges = Int32[]
                     values = Any[]
-                    sizehint!(edges, length(phi.edges)); sizehint!(values, length(renamed_values))
+                    sizehint!(edges, length(phi.edges)); sizehint!(values, length(phi.values))
                     for old_index in 1:length(phi.edges)
                         old_edge = phi.edges[old_index]
                         new_edge = bb_rename_pred[old_edge]
                         if new_edge > 0
                             push!(edges, new_edge)
-                            if isassigned(renamed_values, old_index)
-                                push!(values, renamed_values[old_index])
+                            if isassigned(phi.values, old_index)
+                                val = process_phinode_value(phi.values, old_index, late_fixup, already_inserted, compact.result_idx, ssa_rename, used_ssas, new_new_used_ssas, true, nothing)
+                                push!(values, val)
                             else
                                 resize!(values, length(values)+1)
                             end
                         elseif new_edge == -1
                             @assert length(phi.edges) == 1
-                            if isassigned(renamed_values, old_index)
+                            if isassigned(phi.values, old_index)
+                                val = process_phinode_value(phi.values, old_index, late_fixup, already_inserted, compact.result_idx, ssa_rename, used_ssas, new_new_used_ssas, true, nothing)
                                 push!(edges, -1)
-                                push!(values, renamed_values[old_index])
+                                push!(values, val)
                             end
                         elseif new_edge == -3
                             # Multiple predecessors, we need to expand out this phi
                             all_new_preds = Int32[]
                             add_preds!(all_new_preds, bbs, bb_rename_pred, old_edge)
                             append!(edges, all_new_preds)
-                            if isassigned(renamed_values, old_index)
-                                val = renamed_values[old_index]
-                                for _ in 1:length(all_new_preds)
-                                    push!(values, val)
+                            np = length(all_new_preds)
+                            if np > 0
+                                if isassigned(phi.values, old_index)
+                                    val = process_phinode_value(phi.values, old_index, late_fixup, already_inserted, compact.result_idx, ssa_rename, used_ssas, new_new_used_ssas, true, nothing)
+                                    for p in 1:np
+                                        push!(values, val)
+                                        p > 2 && count_added_node!(compact, val)
+                                    end
+                                else
+                                    resize!(values, length(values)+np)
                                 end
-                                length(all_new_preds) == 0 && kill_current_use!(compact, val)
-                                for _ in 2:length(all_new_preds)
-                                    count_added_node!(compact, val)
-                                end
-                            else
-                                resize!(values, length(values)+length(all_new_preds))
                             end
-                        else
-                            isassigned(renamed_values, old_index) && kill_current_use!(compact, renamed_values[old_index])
                         end
                     end
                     if length(edges) == 0 || (length(edges) == 1 && !isassigned(values, 1))

src/typeinfer.jl

@@ -93,19 +93,21 @@ If set to `true`, record per-method-instance timings within type inference in th
 __set_measure_typeinf(onoff::Bool) = __measure_typeinf__[] = onoff
 const __measure_typeinf__ = RefValue{Bool}(false)
 
-function finish!(interp::AbstractInterpreter, caller::InferenceState, validation_world::UInt)
+function result_edges(interp::AbstractInterpreter, caller::InferenceState)
     result = caller.result
     opt = result.src
-    if opt isa OptimizationState
-        src = ir_to_codeinf!(opt)
-        edges = src.edges::SimpleVector
-        caller.src = result.src = src
+    if isa(opt, OptimizationState)
+        return Core.svec(opt.inlining.edges...)
     else
-        edges = Core.svec(caller.edges...)
-        caller.src.edges = edges
+        return Core.svec(caller.edges...)
     end
+end
+
+function finish!(interp::AbstractInterpreter, caller::InferenceState, validation_world::UInt)
+    result = caller.result
     #@assert last(result.valid_worlds) <= get_world_counter() || isempty(caller.edges)
     if isdefined(result, :ci)
+        edges = result_edges(interp, caller)
         ci = result.ci
         # if we aren't cached, we don't need this edge
         # but our caller might, so let's just make it anyways
@@ -119,9 +121,10 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState, validation
         const_flag = is_result_constabi_eligible(result)
         discard_src = caller.cache_mode === CACHE_MODE_NULL || const_flag
         if !discard_src
-            inferred_result = transform_result_for_cache(interp, result)
+            inferred_result = transform_result_for_cache(interp, result, edges)
             # TODO: do we want to augment edges here with any :invoke targets that we got from inlining (such that we didn't have a direct edge to it already)?
             if inferred_result isa CodeInfo
+                result.src = inferred_result
                 if may_compress(interp)
                     nslots = length(inferred_result.slotflags)
                     resize!(inferred_result.slottypes::Vector{Any}, nslots)
@@ -278,7 +281,16 @@ function is_result_constabi_eligible(result::InferenceResult)
     return isa(result_type, Const) && is_foldable_nothrow(result.ipo_effects) && is_inlineable_constant(result_type.val)
 end
 
-transform_result_for_cache(::AbstractInterpreter, result::InferenceResult) = result.src
+function transform_result_for_cache(::AbstractInterpreter, result::InferenceResult, edges::SimpleVector)
+    src = result.src
+    if isa(src, OptimizationState)
+        src = ir_to_codeinf!(src)
+    end
+    if isa(src, CodeInfo)
+        src.edges = edges
+    end
+    return src
+end
 
 function maybe_compress_codeinfo(interp::AbstractInterpreter, mi::MethodInstance, ci::CodeInfo)
     def = mi.def
@@ -1064,6 +1076,7 @@ function typeinf_frame(interp::AbstractInterpreter, mi::MethodInstance, run_opti
             opt = OptimizationState(frame, interp)
             optimize(interp, opt, frame.result)
             src = ir_to_codeinf!(opt)
+            src.edges = Core.svec(opt.inlining.edges...)
         end
         result.src = frame.src = src
     end

test/AbstractInterpreter.jl

@@ -14,7 +14,7 @@ Compiler.may_optimize(::AbsIntOnlyInterp1) = false
 # it should work even if the interpreter discards inferred source entirely
 @newinterp AbsIntOnlyInterp2
 Compiler.may_optimize(::AbsIntOnlyInterp2) = false
-Compiler.transform_result_for_cache(::AbsIntOnlyInterp2, ::Compiler.InferenceResult) = nothing
+Compiler.transform_result_for_cache(::AbsIntOnlyInterp2, ::Compiler.InferenceResult, edges::Core.SimpleVector) = nothing
 @test Base.infer_return_type(Base.init_stdio, (Ptr{Cvoid},); interp=AbsIntOnlyInterp2()) >: IO
 
 # OverlayMethodTable
@@ -493,9 +493,9 @@ struct CustomData
     inferred
     CustomData(@nospecialize inferred) = new(inferred)
 end
-function Compiler.transform_result_for_cache(interp::CustomDataInterp, result::Compiler.InferenceResult)
+function Compiler.transform_result_for_cache(interp::CustomDataInterp, result::Compiler.InferenceResult, edges::Core.SimpleVector)
     inferred_result = @invoke Compiler.transform_result_for_cache(
-        interp::Compiler.AbstractInterpreter, result::Compiler.InferenceResult)
+        interp::Compiler.AbstractInterpreter, result::Compiler.InferenceResult, edges::Core.SimpleVector)
     return CustomData(inferred_result)
 end
 function Compiler.src_inlining_policy(interp::CustomDataInterp, @nospecialize(src),