fix some new-edges issues (JuliaLang#56598)

vtjnash · web-flow · commit d9d1fc5be8f4 · 2024-11-20T15:17:40.000-05:00
- incorrect edge types were being added from inlining: there is minimal dispatch info available, so best not to add that (which was already added earlier) as it results in failures to validate later - MethodTable/sig order in edges could confuse the iterator: always put the type before the edge now as that is more consistent - edges wasn't converted to a SimpleVector, so they might get ignored later from being in the wrong format - edges were not populated for optimize=false, which made debugging them more inconvenient Fixes JuliaLang#56577
diff --git a/Compiler/src/optimize.jl b/Compiler/src/optimize.jl
@@ -224,7 +224,7 @@ include("ssair/irinterp.jl")
 function ir_to_codeinf!(opt::OptimizationState)
     (; linfo, src) = opt
     src = ir_to_codeinf!(src, opt.ir::IRCode)
-    src.edges = opt.inlining.edges
+    src.edges = Core.svec(opt.inlining.edges...)
     opt.ir = nothing
     maybe_validate_code(linfo, src, "optimized")
     return src
diff --git a/Compiler/src/ssair/inlining.jl b/Compiler/src/ssair/inlining.jl
@@ -64,20 +64,11 @@ end
 
 struct InliningEdgeTracker
     edges::Vector{Any}
-    invokesig::Union{Nothing,Vector{Any}}
-    InliningEdgeTracker(state::InliningState, invokesig::Union{Nothing,Vector{Any}}=nothing) =
-        new(state.edges, invokesig)
+    InliningEdgeTracker(state::InliningState) = new(state.edges)
 end
 
-function add_inlining_edge!(et::InliningEdgeTracker, edge::Union{CodeInstance,MethodInstance})
-    (; edges, invokesig) = et
-    if invokesig === nothing
-        add_one_edge!(edges, edge)
-    else # invoke backedge
-        add_invoke_edge!(edges, invoke_signature(invokesig), edge)
-    end
-    return nothing
-end
+add_inlining_edge!(et::InliningEdgeTracker, edge::CodeInstance) = add_inlining_edge!(et.edges, edge)
+add_inlining_edge!(et::InliningEdgeTracker, edge::MethodInstance) = add_inlining_edge!(et.edges, edge)
 
 function ssa_inlining_pass!(ir::IRCode, state::InliningState, propagate_inbounds::Bool)
     # Go through the function, performing simple inlining (e.g. replacing call by constants
@@ -795,10 +786,7 @@ function compileable_specialization(mi::MethodInstance, effects::Effects,
             return nothing
         end
     end
-    add_inlining_edge!(et, mi) # to the dispatch lookup
-    if mi_invoke !== mi
-        add_invoke_edge!(et.edges, method.sig, mi_invoke) # add_inlining_edge to the invoke call, if that is different
-    end
+    add_inlining_edge!(et, mi_invoke) # to the dispatch lookup
     return InvokeCase(mi_invoke, effects, info)
 end
 
@@ -834,9 +822,8 @@ end
 
 # the general resolver for usual and const-prop'ed calls
 function resolve_todo(mi::MethodInstance, result::Union{Nothing,InferenceResult,VolatileInferenceResult},
-    @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
-    invokesig::Union{Nothing,Vector{Any}}=nothing)
-    et = InliningEdgeTracker(state, invokesig)
+    @nospecialize(info::CallInfo), flag::UInt32, state::InliningState)
+    et = InliningEdgeTracker(state)
 
     preserve_local_sources = true
     if isa(result, InferenceResult)
@@ -922,7 +909,7 @@ end
 
 function analyze_method!(match::MethodMatch, argtypes::Vector{Any},
     @nospecialize(info::CallInfo), flag::UInt32, state::InliningState;
-    allow_typevars::Bool, invokesig::Union{Nothing,Vector{Any}}=nothing,
+    allow_typevars::Bool,
     volatile_inf_result::Union{Nothing,VolatileInferenceResult}=nothing)
     method = match.method
 
@@ -953,7 +940,7 @@ function analyze_method!(match::MethodMatch, argtypes::Vector{Any},
     # Get the specialization for this method signature
     # (later we will decide what to do with it)
     mi = specialize_method(match)
-    return resolve_todo(mi, volatile_inf_result, info, flag, state; invokesig)
+    return resolve_todo(mi, volatile_inf_result, info, flag, state)
 end
 
 function retrieve_ir_for_inlining(cached_result::CodeInstance, src::String)
@@ -1164,9 +1151,8 @@ function handle_invoke_call!(todo::Vector{Pair{Int,Any}},
         return nothing
     end
     result = info.result
-    invokesig = sig.argtypes
     if isa(result, ConcreteResult)
-        item = concrete_result_item(result, info, state; invokesig)
+        item = concrete_result_item(result, info, state)
     elseif isa(result, SemiConcreteResult)
         item = semiconcrete_result_item(result, info, flag, state)
     else
@@ -1175,13 +1161,13 @@ function handle_invoke_call!(todo::Vector{Pair{Int,Any}},
             mi = result.result.linfo
             validate_sparams(mi.sparam_vals) || return nothing
             if Union{} !== argtypes_to_type(argtypes) <: mi.def.sig
-                item = resolve_todo(mi, result.result, info, flag, state; invokesig)
+                item = resolve_todo(mi, result.result, info, flag, state)
                 handle_single_case!(todo, ir, idx, stmt, item, true)
                 return nothing
             end
         end
         volatile_inf_result = result isa VolatileInferenceResult ? result : nothing
-        item = analyze_method!(match, argtypes, info, flag, state; allow_typevars=false, invokesig, volatile_inf_result)
+        item = analyze_method!(match, argtypes, info, flag, state; allow_typevars=false, volatile_inf_result)
     end
     handle_single_case!(todo, ir, idx, stmt, item, true)
     return nothing
@@ -1477,10 +1463,9 @@ end
 may_inline_concrete_result(result::ConcreteResult) =
     isdefined(result, :result) && is_inlineable_constant(result.result)
 
-function concrete_result_item(result::ConcreteResult, @nospecialize(info::CallInfo), state::InliningState;
-    invokesig::Union{Nothing,Vector{Any}}=nothing)
+function concrete_result_item(result::ConcreteResult, @nospecialize(info::CallInfo), state::InliningState)
     if !may_inline_concrete_result(result)
-        et = InliningEdgeTracker(state, invokesig)
+        et = InliningEdgeTracker(state)
         return compileable_specialization(result.edge.def, result.effects, et, info, state)
     end
     @assert result.effects === EFFECTS_TOTAL
diff --git a/Compiler/src/stmtinfo.jl b/Compiler/src/stmtinfo.jl
@@ -49,14 +49,15 @@ function _add_edges_impl(edges::Vector{Any}, info::MethodMatchInfo, mi_edge::Boo
     if !fully_covering(info)
         # add legacy-style missing backedge info also
         exists = false
-        for i in 1:length(edges)
-            if edges[i] === info.mt && edges[i+1] == info.atype
+        for i in 2:length(edges)
+            if edges[i] === info.mt && edges[i-1] == info.atype
                 exists = true
                 break
             end
         end
         if !exists
-            push!(edges, info.mt, info.atype)
+            push!(edges, info.atype)
+            push!(edges, info.mt)
         end
     end
     nmatches = length(info.results)
@@ -98,22 +99,27 @@ function _add_edges_impl(edges::Vector{Any}, info::MethodMatchInfo, mi_edge::Boo
     nothing
 end
 function add_one_edge!(edges::Vector{Any}, edge::MethodInstance)
-    for i in 1:length(edges)
+    i = 1
+    while i <= length(edges)
         edgeᵢ = edges[i]
+        edgeᵢ isa Int && (i += 2 + edgeᵢ; continue)
         edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def)
-        edgeᵢ isa MethodInstance || continue
+        edgeᵢ isa MethodInstance || (i += 1; continue)
         if edgeᵢ === edge && !(i > 1 && edges[i-1] isa Type)
             return # found existing covered edge
         end
+        i += 1
     end
     push!(edges, edge)
     nothing
 end
 function add_one_edge!(edges::Vector{Any}, edge::CodeInstance)
-    for i in 1:length(edges)
+    i = 1
+    while i <= length(edges)
         edgeᵢ_orig = edgeᵢ = edges[i]
+        edgeᵢ isa Int && (i += 2 + edgeᵢ; continue)
         edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def)
-        edgeᵢ isa MethodInstance || continue
+        edgeᵢ isa MethodInstance || (i += 1; continue)
         if edgeᵢ === edge.def && !(i > 1 && edges[i-1] isa Type)
             if edgeᵢ_orig isa MethodInstance
                 # found edge we can upgrade
@@ -123,6 +129,7 @@ function add_one_edge!(edges::Vector{Any}, edge::CodeInstance)
                 return
             end
         end
+        i += 1
     end
     push!(edges, edge)
     nothing
@@ -296,7 +303,8 @@ function add_invoke_edge!(edges::Vector{Any}, @nospecialize(atype), edge::Union{
             end
         end
     end
-    push!(edges, atype, edge)
+    push!(edges, atype)
+    push!(edges, edge)
     nothing
 end
 function add_invoke_edge!(edges::Vector{Any}, @nospecialize(atype), edge::CodeInstance)
@@ -317,10 +325,61 @@ function add_invoke_edge!(edges::Vector{Any}, @nospecialize(atype), edge::CodeIn
             end
         end
     end
-    push!(edges, atype, edge)
+    push!(edges, atype)
+    push!(edges, edge)
     nothing
 end
 
+function add_inlining_edge!(edges::Vector{Any}, edge::MethodInstance)
+    # check if we already have an edge to this code
+    i = 1
+    while i <= length(edges)
+        edgeᵢ = edges[i]
+        if edgeᵢ isa Method && edgeᵢ === edge.def
+            # found edge we can upgrade
+            edges[i] = edge
+            return
+        end
+        edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def)
+        if edgeᵢ isa MethodInstance && edgeᵢ === edge
+            return # found existing covered edge
+        end
+        i += 1
+    end
+    # add_invoke_edge alone
+    push!(edges, (edge.def::Method).sig)
+    push!(edges, edge)
+    nothing
+end
+function add_inlining_edge!(edges::Vector{Any}, edge::CodeInstance)
+    # check if we already have an edge to this code
+    i = 1
+    while i <= length(edges)
+        edgeᵢ = edges[i]
+        if edgeᵢ isa Method && edgeᵢ === edge.def.def
+            # found edge we can upgrade
+            edges[i] = edge
+            return
+        end
+        if edgeᵢ isa MethodInstance && edgeᵢ === edge.def
+            # found edge we can upgrade
+            edges[i] = edge
+            return
+        end
+        if edgeᵢ isa CodeInstance && edgeᵢ.def === edge.def
+            # found existing edge
+            # XXX compare `CodeInstance` identify?
+            return
+        end
+        i += 1
+    end
+    # add_invoke_edge alone
+    push!(edges, (edge.def.def::Method).sig)
+    push!(edges, edge)
+    nothing
+end
+
+
 """
     info::OpaqueClosureCallInfo
 
diff --git a/Compiler/src/typeinfer.jl b/Compiler/src/typeinfer.jl
@@ -97,7 +97,12 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState;
     result = caller.result
     opt = result.src
     if opt isa OptimizationState
-        result.src = ir_to_codeinf!(opt)
+        src = ir_to_codeinf!(opt)
+        edges = src.edges::SimpleVector
+        caller.src = result.src = src
+    else
+        edges = Core.svec(caller.edges...)
+        caller.src.edges = edges
     end
     #@assert last(result.valid_worlds) <= get_world_counter() || isempty(caller.edges)
     if isdefined(result, :ci)
@@ -112,7 +117,7 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState;
         if last(result.valid_worlds) == typemax(UInt)
             # if we can record all of the backedges in the global reverse-cache,
             # we can now widen our applicability in the global cache too
-            store_backedges(ci, caller.edges)
+            store_backedges(ci, edges)
         end
         inferred_result = nothing
         relocatability = 0x1
@@ -142,7 +147,7 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState;
         end
         ccall(:jl_update_codeinst, Cvoid, (Any, Any, Int32, UInt, UInt, UInt32, Any, UInt8, Any, Any),
             ci, inferred_result, const_flag, first(result.valid_worlds), last(result.valid_worlds), encode_effects(result.ipo_effects),
-            result.analysis_results, relocatability, di, Core.svec(caller.edges...))
+            result.analysis_results, relocatability, di, edges)
         engine_reject(interp, ci)
     end
     return nothing
@@ -488,14 +493,43 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
 end
 
 # record the backedges
-function store_backedges(caller::CodeInstance, edges::Vector{Any})
+function store_backedges(caller::CodeInstance, edges::SimpleVector)
     isa(caller.def.def, Method) || return # don't add backedges to toplevel method instance
-    for itr in BackedgeIterator(edges)
-        callee = itr.caller
-        if isa(callee, MethodInstance)
-            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), callee, itr.sig, caller)
+    i = 1
+    while true
+        i > length(edges) && return nothing
+        item = edges[i]
+        if item isa Int
+            i += 2
+            continue # ignore the query information if present but process the contents
+        elseif isa(item, Method)
+            # ignore `Method`-edges (from e.g. failed `abstract_call_method`)
+            i += 1
+            continue
+        end
+        if isa(item, CodeInstance)
+            item = item.def
+        end
+        if isa(item, MethodInstance) # regular dispatch
+            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), item, nothing, caller)
+            i += 1
         else
-            ccall(:jl_method_table_add_backedge, Cvoid, (Any, Any, Any), callee, itr.sig, caller)
+            callee = edges[i+1]
+            if isa(callee, MethodTable) # abstract dispatch (legacy style edges)
+                ccall(:jl_method_table_add_backedge, Cvoid, (Any, Any, Any), callee, item, caller)
+                i += 2
+                continue
+            end
+            # `invoke` edge
+            if isa(callee, Method)
+                # ignore `Method`-edges (from e.g. failed `abstract_call_method`)
+                i += 2
+                continue
+            elseif isa(callee, CodeInstance)
+                callee = callee.def
+            end
+            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), callee, item, caller)
+            i += 2
         end
     end
     nothing
@@ -734,13 +768,14 @@ function codeinst_as_edge(interp::AbstractInterpreter, sv::InferenceState)
     if max_world >= get_world_counter()
         max_world = typemax(UInt)
     end
+    edges = Core.svec(sv.edges...)
     ci = CodeInstance(mi, owner, Any, Any, nothing, nothing, zero(Int32),
-        min_world, max_world, zero(UInt32), nothing, zero(UInt8), nothing, Core.svec(sv.edges...))
+        min_world, max_world, zero(UInt32), nothing, zero(UInt8), nothing, edges)
     if max_world == typemax(UInt)
         # if we can record all of the backedges in the global reverse-cache,
         # we can now widen our applicability in the global cache too
         # TODO: this should probably come after we decide this edge is even useful
-        store_backedges(ci, sv.edges)
+        store_backedges(ci, edges)
     end
     return ci
 end
diff --git a/Compiler/src/utilities.jl b/Compiler/src/utilities.jl
diff --git a/src/staticdata_utils.c b/src/staticdata_utils.c
