inference: parameterize some of hard-coded inference logic (JuliaLang…

…#39439)

This commit parameterizes some of hard-coded inference logic:
- to bail out from inference when a lattice element can't be refined or
  a current inference frame is proven to throw or to be dead
- to add call backedges when the call return type won't be refined

Those `AbstractInterpreter`s used for code optimization (including
`NativeInterpreter`) usually just want the methods defined for
`AbstractInterpreter`, but some other `AbstractInterpreter` may want
other implementations and heuristics to control inference process.
For example, [`JETInterpreter`](https://github.com/aviatesk/JET.jl) is
used for code analysis and wants to add call backedges even when a call
return type is `Any`.

base/compiler/abstractinterpretation.jl

@@ -97,10 +97,9 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
  napplicable = length(applicable)
  rettype = Bottom
  edgecycle = false
- edges = Any[]
+ edges = MethodInstance[]
  nonbot = 0 # the index of the only non-Bottom inference result if > 0
  seen = 0 # number of signatures actually inferred
- istoplevel = sv.linfo.def isa Module
  multiple_matches = napplicable > 1
 
  if f !== nothing && napplicable == 1 && is_method_pure(applicable[1]::MethodMatch)
@@ -115,7 +114,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
  match = applicable[i]::MethodMatch
  method = match.method
  sig = match.spec_types
- if istoplevel && !isdispatchtuple(sig)
+ if bail_out_toplevel_call(interp, sig, sv)
  # only infer concrete call sites in top-level expressions
  add_remark!(interp, sv, "Refusing to infer non-concrete call site in top-level expression")
  rettype = Any
@@ -135,7 +134,9 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
  end
  edgecycle |= edgecycle1::Bool
  this_rt = tmerge(this_rt, rt)
- this_rt === Any && break
+ if bail_out_call(interp, this_rt, sv)
+ break
+ end
  end
  else
  this_rt, edgecycle1, edge = abstract_call_method(interp, method, sig, match.sparams, multiple_matches, sv)
@@ -153,7 +154,9 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
  end
  seen += 1
  rettype = tmerge(rettype, this_rt)
- rettype === Any && break
+ if bail_out_call(interp, rettype, sv)
+ break
+ end
  end
  # try constant propagation if only 1 method is inferred to non-Bottom
  # this is in preparation for inlining, or improving the return result
@@ -179,18 +182,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
  # and avoid keeping track of a more complex result type.
  rettype = Any
  end
- if !(rettype === Any) # adding a new method couldn't refine (widen) this type
- for edge in edges
- add_backedge!(edge::MethodInstance, sv)
- end
- for (thisfullmatch, mt) in zip(fullmatch, mts)
- if !thisfullmatch
- # also need an edge to the method table in case something gets
- # added that did not intersect with any existing method
- add_mt_backedge!(mt, atype, sv)
- end
- end
- end
+ add_call_backedges!(interp, rettype, edges, fullmatch, mts, atype, sv)
  #print("=> ", rettype, "\n")
  if rettype isa LimitedAccuracy
  union!(sv.pclimitations, rettype.causes)
@@ -205,6 +197,27 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
  return CallMeta(rettype, info)
 end
 
+function add_call_backedges!(interp::AbstractInterpreter,
+ @nospecialize(rettype),
+ edges::Vector{MethodInstance},
+ fullmatch::Vector{Bool}, mts::Vector{Core.MethodTable}, @nospecialize(atype),
+ sv::InferenceState)
+ if rettype === Any
+ # for `NativeInterpreter`, we don't add backedges when a new method couldn't refine
+ # (widen) this type
+ return
+ end
+ for edge in edges
+ add_backedge!(edge, sv)
+ end
+ for (thisfullmatch, mt) in zip(fullmatch, mts)
+ if !thisfullmatch
+ # also need an edge to the method table in case something gets
+ # added that did not intersect with any existing method
+ add_mt_backedge!(mt, atype, sv)
+ end
+ end
+end
 
 function const_prop_profitable(@nospecialize(arg))
  # have new information from argtypes that wasn't available from the signature
@@ -746,7 +759,7 @@ function abstract_apply(interp::AbstractInterpreter, @nospecialize(itft), @nospe
  call = abstract_call(interp, nothing, ct, sv, max_methods)
  push!(retinfos, ApplyCallInfo(call.info, arginfo))
  res = tmerge(res, call.rt)
- if res === Any
+ if bail_out_apply(interp, res, sv)
  # No point carrying forward the info, we're not gonna inline it anyway
  retinfo = nothing
  break
@@ -1171,7 +1184,7 @@ function abstract_eval_statement(interp::AbstractInterpreter, @nospecialize(e),
  argtypes = Vector{Any}(undef, n)
  @inbounds for i = 1:n
  ai = abstract_eval_value(interp, ea[i], vtypes, sv)
- if ai === Bottom
+ if bail_out_statement(interp, ai, sv)
  return Bottom
  end
  argtypes[i] = ai
@@ -1349,6 +1362,8 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
  condt = abstract_eval_value(interp, stmt.cond, s[pc], frame)
  if condt === Bottom
  empty!(frame.pclimitations)
+ end
+ if bail_out_local(interp, condt, frame)
  break
  end
  condval = maybe_extract_const_bool(condt)
@@ -1440,7 +1455,9 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
  else
  if hd === :(=)
  t = abstract_eval_statement(interp, stmt.args[2], changes, frame)
- t === Bottom && break
+ if bail_out_local(interp, t, frame)
+ break
+ end
  frame.src.ssavaluetypes[pc] = t
  lhs = stmt.args[1]
  if isa(lhs, Slot)
@@ -1455,7 +1472,9 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
  # these do not generate code
  else
  t = abstract_eval_statement(interp, stmt, changes, frame)
- t === Bottom && break
+ if bail_out_local(interp, t, frame)
+ break
+ end
  if !isempty(frame.ssavalue_uses[pc])
  record_ssa_assign(pc, t, frame)
  else

base/compiler/types.jl

@@ -212,3 +212,16 @@ may_compress(ni::NativeInterpreter) = true
 may_discard_trees(ni::NativeInterpreter) = true
 
 method_table(ai::AbstractInterpreter) = InternalMethodTable(get_world_counter(ai))
+
+# define inference bail out logic
+# `NativeInterpreter` bails out from inference when
+# - a lattice element grows up to `Any` (inter-procedural call, abstract apply)
+# - a lattice element gets down to `Bottom` (statement inference, local frame inference)
+# - inferring non-concrete toplevel call sites
+bail_out_call(interp::AbstractInterpreter, @nospecialize(t), sv) = t === Any
+bail_out_apply(interp::AbstractInterpreter, @nospecialize(t), sv) = t === Any
+bail_out_statement(interp::AbstractInterpreter, @nospecialize(t), sv) = t === Bottom
+bail_out_local(interp::AbstractInterpreter, @nospecialize(t), sv) = t === Bottom
+function bail_out_toplevel_call(interp::AbstractInterpreter, @nospecialize(sig), sv)
+ return isa(sv.linfo.def, Module) && !isdispatchtuple(sig)
+end