Skip to content

Commit

Permalink
convert to full CodeInstance edge list on the Julia side
Browse files Browse the repository at this point in the history
  • Loading branch information
aviatesk committed Oct 25, 2024
1 parent 6b1b0c1 commit 140f560
Show file tree
Hide file tree
Showing 7 changed files with 90 additions and 70 deletions.
83 changes: 41 additions & 42 deletions base/compiler/ssair/inlining.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ end

struct ConstantCase
val::Any
ConstantCase(@nospecialize val) = new(val)
edge::CodeInstance
ConstantCase(@nospecialize(val), edge::CodeInstance) = new(val, edge)
end

struct SomeCase
Expand Down Expand Up @@ -68,12 +69,12 @@ struct InliningEdgeTracker
new(state.edges, invokesig)
end

function add_inlining_edge!(et::InliningEdgeTracker, mi::MethodInstance)
function add_inlining_edge!(et::InliningEdgeTracker, edge::CodeInstance)
(; edges, invokesig) = et
if invokesig === nothing
add_one_edge!(edges, mi)
add_one_edge!(edges, edge)
else # invoke backedge
add_invoke_edge!(edges, invoke_signature(invokesig), mi)
add_invoke_edge!(edges, invoke_signature(invokesig), edge)
end
return nothing
end
Expand Down Expand Up @@ -784,11 +785,11 @@ function rewrite_apply_exprargs!(todo::Vector{Pair{Int,Any}},
return new_argtypes
end

function compileable_specialization(mi::MethodInstance, effects::Effects,
et::InliningEdgeTracker, @nospecialize(info::CallInfo); compilesig_invokes::Bool=true)
mi_invoke = mi
function compileable_specialization(edge::CodeInstance, effects::Effects,
et::InliningEdgeTracker, @nospecialize(info::CallInfo), state::InliningState)
mi_invoke = mi = edge.def
method, atype, sparams = mi.def::Method, mi.specTypes, mi.sparam_vals
if compilesig_invokes
if OptimizationParams(state.interp).compilesig_invokes
new_atype = get_compileable_sig(method, atype, sparams)
new_atype === nothing && return nothing
if atype !== new_atype
Expand All @@ -806,43 +807,42 @@ function compileable_specialization(mi::MethodInstance, effects::Effects,
return nothing
end
end
add_inlining_edge!(et, mi) # to the dispatch lookup
mi_invoke !== mi && add_invoke_edge!(et.edges, method.sig, mi_invoke) # add_inlining_edge to the invoke call, if that is different
add_inlining_edge!(et, edge) # to the dispatch lookup
if mi_invoke !== mi
edge_invoke = codeinst_as_invoke_edge(state.interp, mi_invoke)
add_invoke_edge!(et.edges, method.sig, edge_invoke) # add_inlining_edge to the invoke call, if that is different
end
return InvokeCase(mi_invoke, effects, info)
end

function compileable_specialization(match::MethodMatch, effects::Effects,
et::InliningEdgeTracker, @nospecialize(info::CallInfo); compilesig_invokes::Bool=true)
mi = specialize_method(match)
return compileable_specialization(mi, effects, et, info; compilesig_invokes)
end

struct InferredResult
src::Any # CodeInfo or IRCode
effects::Effects
InferredResult(@nospecialize(src), effects::Effects) = new(src, effects)
edge::CodeInstance
InferredResult(@nospecialize(src), effects::Effects, edge::CodeInstance) = new(src, effects, edge)
end
@inline function get_cached_result(state::InliningState, mi::MethodInstance)
code = get(code_cache(state), mi, nothing)
if code isa CodeInstance
if use_const_api(code)
# in this case function can be inlined to a constant
return ConstantCase(quoted(code.rettype_const))
return ConstantCase(quoted(code.rettype_const), code)
end
return code
end
return nothing
end
@inline function get_local_result(inf_result::InferenceResult)
@assert isdefined(inf_result, :ci_as_edge) "InferenceResult without ci_as_edge"
effects = inf_result.ipo_effects
if is_foldable_nothrow(effects)
res = inf_result.result
if isa(res, Const) && is_inlineable_constant(res.val)
# use constant calling convention
return ConstantCase(quoted(res.val))
return ConstantCase(quoted(res.val), inf_result.ci_as_edge)
end
end
return InferredResult(inf_result.src, effects)
return InferredResult(inf_result.src, effects, inf_result.ci_as_edge)
end

# the general resolver for usual and const-prop'ed calls
Expand All @@ -862,30 +862,29 @@ function resolve_todo(mi::MethodInstance, result::Union{Nothing,InferenceResult,
inferred_result = get_cached_result(state, mi)
end
if inferred_result isa ConstantCase
add_inlining_edge!(et, mi)
add_inlining_edge!(et, inferred_result.edge)
return inferred_result
elseif inferred_result isa InferredResult
(; src, effects) = inferred_result
(; src, effects, edge) = inferred_result
elseif inferred_result isa CodeInstance
src = @atomic :monotonic inferred_result.inferred
effects = decode_effects(inferred_result.ipo_purity_bits)
edge = inferred_result
else # there is no cached source available, bail out
return compileable_specialization(mi, Effects(), et, info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
edge = codeinst_as_invoke_edge(state.interp, mi)
return compileable_specialization(edge, Effects(), et, info, state)
end

# the duplicated check might have been done already within `analyze_method!`, but still
# we need it here too since we may come here directly using a constant-prop' result
if !OptimizationParams(state.interp).inlining || is_stmt_noinline(flag)
return compileable_specialization(mi, effects, et, info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
return compileable_specialization(edge, effects, et, info, state)
end

src_inlining_policy(state.interp, src, info, flag) ||
return compileable_specialization(mi, effects, et, info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
return compileable_specialization(edge, effects, et, info, state)

add_inlining_edge!(et, mi)
add_inlining_edge!(et, edge)
if inferred_result isa CodeInstance
ir, spec_info, debuginfo = retrieve_ir_for_inlining(inferred_result, src)
else
Expand All @@ -905,7 +904,7 @@ function resolve_todo(mi::MethodInstance, @nospecialize(info::CallInfo), flag::U

cached_result = get_cached_result(state, mi)
if cached_result isa ConstantCase
add_inlining_edge!(et, mi)
add_inlining_edge!(et, cached_result.edge)
return cached_result
elseif cached_result isa CodeInstance
src = @atomic :monotonic cached_result.inferred
Expand All @@ -916,7 +915,7 @@ function resolve_todo(mi::MethodInstance, @nospecialize(info::CallInfo), flag::U

src_inlining_policy(state.interp, src, info, flag) || return nothing
ir, spec_info, debuginfo = retrieve_ir_for_inlining(cached_result, src)
add_inlining_edge!(et, mi)
add_inlining_edge!(et, cached_result)
return InliningTodo(mi, ir, spec_info, debuginfo, effects)
end

Expand Down Expand Up @@ -1464,14 +1463,12 @@ function semiconcrete_result_item(result::SemiConcreteResult,
# a `@noinline`-declared method when it's marked as `@constprop :aggressive`.
# Suppress the inlining here (unless inlining is requested at the callsite).
(is_declared_noinline(mi.def::Method) && !is_stmt_inline(flag)))
return compileable_specialization(mi, result.effects, et, info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
return compileable_specialization(result.edge, result.effects, et, info, state)
end
src_inlining_policy(state.interp, result.ir, info, flag) ||
return compileable_specialization(mi, result.effects, et, info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
return compileable_specialization(result.edge, result.effects, et, info, state)

add_inlining_edge!(et, mi)
add_inlining_edge!(et, result.edge)
preserve_local_sources = OptimizationParams(state.interp).preserve_local_sources
ir, _, debuginfo = retrieve_ir_for_inlining(mi, result.ir, preserve_local_sources)
return InliningTodo(mi, ir, result.spec_info, debuginfo, result.effects)
Expand Down Expand Up @@ -1503,11 +1500,10 @@ function concrete_result_item(result::ConcreteResult, @nospecialize(info::CallIn
invokesig::Union{Nothing,Vector{Any}}=nothing)
if !may_inline_concrete_result(result)
et = InliningEdgeTracker(state, invokesig)
return compileable_specialization(result.edge.def, result.effects, et, info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
return compileable_specialization(result.edge, result.effects, et, info, state)
end
@assert result.effects === EFFECTS_TOTAL
return ConstantCase(quoted(result.result))
return ConstantCase(quoted(result.result), result.edge)
end

function handle_cases!(todo::Vector{Pair{Int,Any}}, ir::IRCode, idx::Int, stmt::Expr,
Expand Down Expand Up @@ -1556,11 +1552,14 @@ function handle_modifyop!_call!(ir::IRCode, idx::Int, stmt::Expr, info::ModifyOp
info isa MethodResultPure && (info = info.info)
info isa ConstCallInfo && (info = info.call)
info isa MethodMatchInfo || return nothing
length(info.results) == 1 || return nothing
length(info.edges) == length(info.results) == 1 || return nothing
match = info.results[1]::MethodMatch
match.fully_covers || return nothing
case = compileable_specialization(match, Effects(), InliningEdgeTracker(state), info;
compilesig_invokes=OptimizationParams(state.interp).compilesig_invokes)
edge = info.edges[1]
if edge === nothing
edge = codeinst_as_invoke_edge(state.interp, specialize_method(match))
end
case = compileable_specialization(edge, Effects(), InliningEdgeTracker(state), info, state)
case === nothing && return nothing
stmt.head = :invoke_modify
pushfirst!(stmt.args, case.invoke)
Expand Down
4 changes: 2 additions & 2 deletions base/compiler/ssair/passes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1529,7 +1529,7 @@ function try_inline_finalizer!(ir::IRCode, argexprs::Vector{Any}, idx::Int,
if code isa CodeInstance
if use_const_api(code)
# No code in the function - Nothing to do
add_inlining_edge!(et, mi)
add_inlining_edge!(et, code)
return true
end
src = @atomic :monotonic code.inferred
Expand All @@ -1544,7 +1544,7 @@ function try_inline_finalizer!(ir::IRCode, argexprs::Vector{Any}, idx::Int,
length(src.cfg.blocks) == 1 || return false

# Ok, we're committed to inlining the finalizer
add_inlining_edge!(et, mi)
add_inlining_edge!(et, code)

# TODO: Should there be a special line number node for inlined finalizers?
inline_at = ir[SSAValue(idx)][:line]
Expand Down
24 changes: 14 additions & 10 deletions base/compiler/stmtinfo.jl
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ function add_edges_impl(edges::Vector{Any}, info::MethodMatchInfo)
# if this doesn't succeed, the backedge will be less precise,
# but the forward edge will maintain the precision
if edge.def.specTypes === info.results[1].spec_types
add_one_edge!(edges, edge.def) # TODO CodeInstance as edge
add_one_edge!(edges, edge)
return nothing
end
end
Expand All @@ -81,18 +81,20 @@ function add_edges_impl(edges::Vector{Any}, info::MethodMatchInfo)
for i = 1:nmatches
edge = info.edges[i]
if edge !== nothing
push!(edges, edge.def) # TODO CodeInstance as edge
push!(edges, edge)
end
end
nothing
end
function add_one_edge!(edges::Vector{Any}, mi::MethodInstance)
function add_one_edge!(edges::Vector{Any}, edge::CodeInstance)
for i in 1:length(edges)
if edges[i] === mi && !(i > 1 && edges[i-1] isa Type)
edgeᵢ = edges[i]
# XXX compare `CodeInstance` identify?
if edgeᵢ isa CodeInstance && edgeᵢ.def === edge.def && !(i > 1 && edges[i-1] isa Type)
return
end
end
push!(edges, mi)
push!(edges, edge)
nothing
end
nsplit_impl(info::MethodMatchInfo) = 1
Expand Down Expand Up @@ -264,19 +266,21 @@ end
function add_edges_impl(edges::Vector{Any}, info::InvokeCallInfo)
edge = info.edge
if edge !== nothing
add_invoke_edge!(edges, info.atype, edge.def) # TODO CodeInstance as edge
add_invoke_edge!(edges, info.atype, edge)
end
nothing
end
function add_invoke_edge!(edges::Vector{Any}, @nospecialize(atype), mi::MethodInstance)
function add_invoke_edge!(edges::Vector{Any}, @nospecialize(atype), edge::CodeInstance)
for i in 2:length(edges)
if edges[i] === mi
edgeᵢ = edges[i]
if edgeᵢ isa CodeInstance && edgeᵢ.def === edge.def # XXX compare `CodeInstance` identify?
edge_minus_1 = edges[i-1]
if edge_minus_1 isa Type && edge_minus_1 == atype
return nothing
end
end
end
push!(edges, atype, mi)
push!(edges, atype, edge)
nothing
end

Expand All @@ -295,7 +299,7 @@ end
function add_edges_impl(edges::Vector{Any}, info::OpaqueClosureCallInfo)
edge = info.edge
if edge !== nothing
add_one_edge!(edges, edge.def) # TODO CodeInstance as edge
add_one_edge!(edges, edge)
end
end

Expand Down
6 changes: 3 additions & 3 deletions base/compiler/tfuncs.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1432,7 +1432,7 @@ end
return Future{CallMeta}(callinfo, interp, sv) do callinfo, interp, sv
TF = TF.contents
RT = RT.contents
TF2 = tmeet(callinfo.rt, widenconst(TF))
TF2 = tmeet(ipo_lattice(interp), callinfo.rt, widenconst(TF))
if TF2 === Bottom
RT = Bottom
elseif isconcretetype(RT) && has_nontrivial_extended_info(𝕃ᵢ, TF2) # isconcrete condition required to form a PartialStruct
Expand Down Expand Up @@ -3012,7 +3012,7 @@ function abstract_applicable(interp::AbstractInterpreter, argtypes::Vector{Any},
if rt !== Bool
for i = 1:napplicable
(; match, edges, edge_idx) = applicable[i]
edges[edge_idx] = codeinst_as_method_match_edge(interp, specialize_method(match))
edges[edge_idx] = codeinst_as_invoke_edge(interp, specialize_method(match))
end
info = VirtualMethodMatchInfo(matches.info)
end
Expand Down Expand Up @@ -3055,7 +3055,7 @@ function _hasmethod_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, sv
vinfo = MethodMatchInfo(vresults, mt, types, false) # XXX: this should actually be an info with invoke-type edge
else
rt = Const(true)
edge = codeinst_as_method_match_edge(interp, specialize_method(match))
edge = codeinst_as_invoke_edge(interp, specialize_method(match))
vinfo = InvokeCallInfo(edge, match, nothing, types)
end
info = VirtualMethodMatchInfo(vinfo)
Expand Down
Loading

0 comments on commit 140f560

Please sign in to comment.