From 09a27b3fe405ea19344881a07d1e1e60cc11553d Mon Sep 17 00:00:00 2001 From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com> Date: Wed, 14 Feb 2024 19:17:39 +0900 Subject: [PATCH] AbsInt: add interfaces to customize cases when cached results are used (#53318) For external abstract interpreters like JET, which propagate custom data interprocedurally, it is essential to inject custom behaviors into where cached regular/constant inference results are used. Previously, this was accomplished by overloading both `abstract_call_method` and `get(::WorldView{CustomView}, ...)` (and `const_prop_call` and `cached_lookup`), that method was admittedly hacky and should probably better to be avoided. Moreover, after #52233, doing so has become infeasible when the external abstract interpreter uses `InternalCodeCache`. To address this issue, this commit provides an interface named `return_cached_result`. This allows external abstract interpreters to inject custom interprocedural data propagation during abstract interpretation even when they use `InternalCodeCache`. --- base/compiler/abstractinterpretation.jl | 62 ++++++++++++++----------- base/compiler/typeinfer.jl | 20 +++++--- 2 files changed, 48 insertions(+), 34 deletions(-) diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl index 9244ee5f3f8ec..21303f527e3f2 100644 --- a/base/compiler/abstractinterpretation.jl +++ b/base/compiler/abstractinterpretation.jl @@ -1221,45 +1221,53 @@ function semi_concrete_eval_call(interp::AbstractInterpreter, return nothing end +const_prop_result(inf_result::InferenceResult) = + ConstCallResults(inf_result.result, inf_result.exc_result, ConstPropResult(inf_result), + inf_result.ipo_effects, inf_result.linfo) + +# return cached constant analysis result +return_cached_result(::AbstractInterpreter, inf_result::InferenceResult, ::AbsIntState) = + const_prop_result(inf_result) + function const_prop_call(interp::AbstractInterpreter, mi::MethodInstance, result::MethodCallResult, arginfo::ArgInfo, sv::AbsIntState, concrete_eval_result::Union{Nothing, ConstCallResults}=nothing) inf_cache = get_inference_cache(interp) 𝕃ᵢ = typeinf_lattice(interp) inf_result = cache_lookup(𝕃ᵢ, mi, arginfo.argtypes, inf_cache) - if inf_result === nothing - # fresh constant prop' - argtypes = has_conditional(𝕃ᵢ, sv) ? ConditionalArgtypes(arginfo, sv) : SimpleArgtypes(arginfo.argtypes) - inf_result = InferenceResult(mi, argtypes, typeinf_lattice(interp)) - if !any(inf_result.overridden_by_const) - add_remark!(interp, sv, "[constprop] Could not handle constant info in matching_cache_argtypes") - return nothing - end - frame = InferenceState(inf_result, #=cache_mode=#:local, interp) - if frame === nothing - add_remark!(interp, sv, "[constprop] Could not retrieve the source") - return nothing # this is probably a bad generated function (unsound), but just ignore it - end - frame.parent = sv - if !typeinf(interp, frame) - add_remark!(interp, sv, "[constprop] Fresh constant inference hit a cycle") - return nothing - end - @assert inf_result.result !== nothing - if concrete_eval_result !== nothing - # override return type and effects with concrete evaluation result if available - inf_result.result = concrete_eval_result.rt - inf_result.ipo_effects = concrete_eval_result.effects - end - else + if inf_result !== nothing # found the cache for this constant prop' if inf_result.result === nothing add_remark!(interp, sv, "[constprop] Found cached constant inference in a cycle") return nothing end + @assert inf_result.linfo === mi "MethodInstance for cached inference result does not match" + return return_cached_result(interp, inf_result, sv) + end + # perform fresh constant prop' + argtypes = has_conditional(𝕃ᵢ, sv) ? ConditionalArgtypes(arginfo, sv) : SimpleArgtypes(arginfo.argtypes) + inf_result = InferenceResult(mi, argtypes, typeinf_lattice(interp)) + if !any(inf_result.overridden_by_const) + add_remark!(interp, sv, "[constprop] Could not handle constant info in matching_cache_argtypes") + return nothing + end + frame = InferenceState(inf_result, #=cache_mode=#:local, interp) + if frame === nothing + add_remark!(interp, sv, "[constprop] Could not retrieve the source") + return nothing # this is probably a bad generated function (unsound), but just ignore it + end + frame.parent = sv + if !typeinf(interp, frame) + add_remark!(interp, sv, "[constprop] Fresh constant inference hit a cycle") + return nothing + end + @assert inf_result.result !== nothing + if concrete_eval_result !== nothing + # override return type and effects with concrete evaluation result if available + inf_result.result = concrete_eval_result.rt + inf_result.ipo_effects = concrete_eval_result.effects end - return ConstCallResults(inf_result.result, inf_result.exc_result, - ConstPropResult(inf_result), inf_result.ipo_effects, mi) + return const_prop_result(inf_result) end # TODO implement MustAlias forwarding diff --git a/base/compiler/typeinfer.jl b/base/compiler/typeinfer.jl index e30c6e5f96fcb..747fd6e3ef705 100644 --- a/base/compiler/typeinfer.jl +++ b/base/compiler/typeinfer.jl @@ -818,23 +818,29 @@ struct EdgeCallResult end end +# return cached regular inference result +function return_cached_result(::AbstractInterpreter, codeinst::CodeInstance, caller::AbsIntState) + rt = cached_return_type(codeinst) + effects = ipo_effects(codeinst) + update_valid_age!(caller, WorldRange(min_world(codeinst), max_world(codeinst))) + return EdgeCallResult(rt, codeinst.exctype, codeinst.def, effects) +end + # compute (and cache) an inferred AST and return the current best estimate of the result type function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize(atype), sparams::SimpleVector, caller::AbsIntState) mi = specialize_method(method, atype, sparams)::MethodInstance - code = get(code_cache(interp), mi, nothing) + codeinst = get(code_cache(interp), mi, nothing) force_inline = is_stmt_inline(get_curr_ssaflag(caller)) - if code isa CodeInstance # return existing rettype if the code is already inferred - inferred = @atomic :monotonic code.inferred + if codeinst isa CodeInstance # return existing rettype if the code is already inferred + inferred = @atomic :monotonic codeinst.inferred if inferred === nothing && force_inline # we already inferred this edge before and decided to discard the inferred code, # nevertheless we re-infer it here again in order to propagate the re-inferred # source to the inliner as a volatile result cache_mode = CACHE_MODE_VOLATILE else - rt = cached_return_type(code) - effects = ipo_effects(code) - update_valid_age!(caller, WorldRange(min_world(code), max_world(code))) - return EdgeCallResult(rt, code.exctype, mi, effects) + @assert codeinst.def === mi "MethodInstance for cached edge does not match" + return return_cached_result(interp, codeinst, caller) end else cache_mode = CACHE_MODE_GLOBAL # cache edge targets globally by default