enable/improve constant propagation through varargs methods (JuliaLan…

…g#26826)

- Store varargs type information in the InferenceResult object, such that the info can be used during inference/optimization

- Hack in a more precise return type for getfield of a vararg tuple. Ideally, we would handle this by teaching inference to track the types of the individual fields of a Tuple, which would make this unnecessary, but until then, this hack is helpful.

- Spoof parents as well as children during recursion limiting, so that higher degree cycles are appropriately spoofed

- A broadcast test marked as broken is now no longer broken, presumably due to the optimizations in this commit

- Fix relationship between depth/mindepth in limit_type_size/is_derived_type. The relationship should have been inverse over the domain in which they overlap, but was not maintained consistently. An example of problematic case was:

t = Tuple{X,X} where X<:Tuple{Tuple{Int64,Vararg{Int64,N} where N},Tuple{Int64,Vararg{Int64,N} where N}}
c = Tuple{X,X} where X<:Tuple{Int64,Vararg{Int64,N} where N}

because is_derived_type was computing the depth of usage rather than the depth of definition. This change thus makes the depth/mindepth calculations more consistent, and causes the limiting heuristic to return strictly wider types than it did before.

- Move the optimizer's "varargs types to tuple type" rewrite to after cache lookup.Inference is populating the InferenceResult cache using the varargs form, so the optimizer needs to do the lookup before writing the atypes in order to avoid cache misses.

Co-authored-by: Jameson Nash <[email protected]>
Co-authored-by: Keno Fischer <[email protected]>

base/compiler/abstractinterpretation.jl

@@ -116,8 +116,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
  method.isva && (nargs -= 1)
  length(argtypes) >= nargs || return Any # probably limit_tuple_type made this non-matching method apparently match
  haveconst = false
- for i in 1:nargs
- a = argtypes[i]
+ for a in argtypes
  if isa(a, Const) && !isdefined(typeof(a.val), :instance) && !(isa(a.val, Type) && issingletontype(a.val))
  # have new information from argtypes that wasn't available from the signature
  haveconst = true
@@ -144,8 +143,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
  tm = _topmod(sv)
  if !istopfunction(tm, f, :getproperty) && !istopfunction(tm, f, :setproperty!)
  # in this case, see if all of the arguments are constants
- for i in 1:nargs
- a = argtypes[i]
+ for a in argtypes
  if !isa(a, Const) && !isconstType(a)
  return Any
  end
@@ -156,6 +154,17 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
  if inf_result === nothing
  inf_result = InferenceResult(code)
  atypes = get_argtypes(inf_result)
+ if method.isva
+ vargs = argtypes[(nargs + 1):end]
+ for i in 1:length(vargs)
+ a = vargs[i]
+ if i > length(inf_result.vargs)
+ push!(inf_result.vargs, a)
+ elseif a isa Const
+ inf_result.vargs[i] = a
+ end
+ end
+ end
  for i in 1:nargs
  a = argtypes[i]
  if a isa Const
@@ -187,7 +196,13 @@ function abstract_call_method(method::Method, @nospecialize(sig), sparams::Simpl
  cyclei = 0
  infstate = sv
  edgecycle = false
+ # The `method_for_inference_heuristics` will expand the given method's generator if
+ # necessary in order to retrieve this field from the generated `CodeInfo`, if it exists.
+ # The other `CodeInfo`s we inspect will already have this field inflated, so we just
+ # access it directly instead (to avoid regeneration).
  method2 = method_for_inference_heuristics(method, sig, sparams, sv.params.world) # Union{Method, Nothing}
+ sv_method2 = sv.src.method_for_inference_limit_heuristics # limit only if user token match
+ sv_method2 isa Method || (sv_method2 = nothing) # Union{Method, Nothing}
  while !(infstate === nothing)
  infstate = infstate::InferenceState
  if method === infstate.linfo.def
@@ -208,7 +223,9 @@ function abstract_call_method(method::Method, @nospecialize(sig), sparams::Simpl
  for parent in infstate.callers_in_cycle
  # check in the cycle list first
  # all items in here are mutual parents of all others
- if parent.linfo.def === sv.linfo.def
+ parent_method2 = parent.src.method_for_inference_limit_heuristics # limit only if user token match
+ parent_method2 isa Method || (parent_method2 = nothing) # Union{Method, Nothing}
+ if parent.linfo.def === sv.linfo.def && sv_method2 === parent_method2
  topmost = infstate
  edgecycle = true
  break
@@ -218,7 +235,9 @@ function abstract_call_method(method::Method, @nospecialize(sig), sparams::Simpl
  # then check the parent link
  if topmost === nothing && parent !== nothing
  parent = parent::InferenceState
- if parent.cached && parent.linfo.def === sv.linfo.def
+ parent_method2 = parent.src.method_for_inference_limit_heuristics # limit only if user token match
+ parent_method2 isa Method || (parent_method2 = nothing) # Union{Method, Nothing}
+ if parent.cached && parent.linfo.def === sv.linfo.def && sv_method2 === parent_method2
  topmost = infstate
  edgecycle = true
  end
@@ -315,8 +334,8 @@ function precise_container_type(@nospecialize(arg), @nospecialize(typ), vtypes::
  end
 
  arg = ssa_def_expr(arg, sv)
- if is_specializable_vararg_slot(arg, sv)
- return Any[rewrap_unionall(p, sv.linfo.specTypes) for p in sv.vararg_type_container.parameters]
+ if is_specializable_vararg_slot(arg, sv.nargs, sv.result.vargs)
+ return sv.result.vargs
  end
 
  tti0 = widenconst(typ)
@@ -482,7 +501,13 @@ function abstract_call(@nospecialize(f), fargs::Union{Tuple{},Vector{Any}}, argt
  tm = _topmod(sv)
  if isa(f, Builtin) || isa(f, IntrinsicFunction)
  rt = builtin_tfunction(f, argtypes[2:end], sv)
- if (rt === Bool || (isa(rt, Const) && isa(rt.val, Bool))) && isa(fargs, Vector{Any})
+ if f === getfield && isa(fargs, Vector{Any}) && length(argtypes) == 3 && isa(argtypes[3], Const) && isa(argtypes[3].val, Int) && argtypes[2] ⊑ Tuple
+ cti = precise_container_type(fargs[2], argtypes[2], vtypes, sv)
+ idx = argtypes[3].val
+ if 1 <= idx <= length(cti)
+ rt = unwrapva(cti[idx])
+ end
+ elseif (rt === Bool || (isa(rt, Const) && isa(rt.val, Bool))) && isa(fargs, Vector{Any})
  # perform very limited back-propagation of type information for `is` and `isa`
  if f === isa
  a = ssa_def_expr(fargs[2], sv)

base/compiler/inferenceresult.jl

@@ -5,6 +5,9 @@ const EMPTY_VECTOR = Vector{Any}()
 mutable struct InferenceResult
  linfo::MethodInstance
  args::Vector{Any}
+ vargs::Vector{Any} # Memoize vararg type info w/Consts here when calling get_argtypes
+ # on the InferenceResult, so that the optimizer can use this info
+ # later during inlining.
  result # ::Type, or InferenceState if WIP
  src::Union{CodeInfo, Nothing} # if inferred copy is available
  function InferenceResult(linfo::MethodInstance)
@@ -13,7 +16,7 @@ mutable struct InferenceResult
  else
  result = linfo.rettype
  end
- return new(linfo, EMPTY_VECTOR, result, nothing)
+ return new(linfo, EMPTY_VECTOR, Any[], result, nothing)
  end
 end
 
@@ -31,7 +34,35 @@ function get_argtypes(result::InferenceResult)
  end
  vararg_type = Tuple
  else
- vararg_type = rewrap(tupleparam_tail(atypes, nargs), linfo.specTypes)
+ laty = length(atypes)
+ if nargs > laty
+ va = atypes[laty]
+ if isvarargtype(va)
+ new_va = rewrap_unionall(unconstrain_vararg_length(va), linfo.specTypes)
+ vararg_type_vec = Any[new_va]
+ vararg_type = Tuple{new_va}
+ else
+ vararg_type_vec = Any[]
+ vararg_type = Tuple{}
+ end
+ else
+ vararg_type_vec = Any[]
+ for p in atypes[nargs:laty]
+ p = isvarargtype(p) ? unconstrain_vararg_length(p) : p
+ push!(vararg_type_vec, rewrap_unionall(p, linfo.specTypes))
+ end
+ vararg_type = tuple_tfunc(Tuple{vararg_type_vec...})
+ for i in 1:length(vararg_type_vec)
+ atyp = vararg_type_vec[i]
+ if isa(atyp, DataType) && isdefined(atyp, :instance)
+ # replace singleton types with their equivalent Const object
+ vararg_type_vec[i] = Const(atyp.instance)
+ elseif isconstType(atyp)
+ vararg_type_vec[i] = Const(atyp.parameters[1])
+ end
+ end
+ end
+ result.vargs = vararg_type_vec
  end
  args[nargs] = vararg_type
  nargs -= 1
@@ -80,19 +111,12 @@ function cache_lookup(code::MethodInstance, argtypes::Vector{Any}, cache::Vector
  for cache_code in cache
  # try to search cache first
  cache_args = cache_code.args
- if cache_code.linfo === code && length(cache_args) >= nargs
+ cache_vargs = cache_code.vargs
+ if cache_code.linfo === code && length(argtypes) === (length(cache_vargs) + nargs)
  cache_match = true
- # verify that the trailing args (va) aren't Const
- for i in (nargs + 1):length(cache_args)
- if isa(cache_args[i], Const)
- cache_match = false
- break
- end
- end
- cache_match || continue
- for i in 1:nargs
+ for i in 1:length(argtypes)
  a = argtypes[i]
- ca = cache_args[i]
+ ca = i <= nargs ? cache_args[i] : cache_vargs[i - nargs]
  # verify that all Const argument types match between the call and cache
  if (isa(a, Const) || isa(ca, Const)) && !(a === ca)
  cache_match = false

base/compiler/inferencestate.jl

@@ -30,7 +30,6 @@ mutable struct InferenceState
  # ssavalue sparsity and restart info
  ssavalue_uses::Vector{BitSet}
  ssavalue_defs::Vector{LineNum}
- vararg_type_container #::Type
 
  backedges::Vector{Tuple{InferenceState, LineNum}} # call-graph backedges connecting from callee to caller
  callers_in_cycle::Vector{InferenceState}
@@ -102,19 +101,11 @@ mutable struct InferenceState
  # initial types
  nslots = length(src.slotnames)
  argtypes = get_argtypes(result)
- vararg_type_container = nothing
  nargs = length(argtypes)
  s_argtypes = VarTable(undef, nslots)
  src.slottypes = Vector{Any}(undef, nslots)
  for i in 1:nslots
  at = (i > nargs) ? Bottom : argtypes[i]
- if !toplevel && linfo.def.isva && i == nargs
- if !(at == Tuple) # would just be a no-op
- vararg_type_container = unwrap_unionall(at)
- vararg_type = tuple_tfunc(vararg_type_container) # returns a Const object, if applicable
- at = rewrap(vararg_type, linfo.specTypes)
- end
- end
  s_argtypes[i] = VarState(at, i > nargs)
  src.slottypes[i] = at
  end
@@ -152,7 +143,7 @@ mutable struct InferenceState
  nargs, s_types, s_edges,
  Union{}, W, 1, n,
  cur_hand, handler_at, n_handlers,
- ssavalue_uses, ssavalue_defs, vararg_type_container,
+ ssavalue_uses, ssavalue_defs,
  Vector{Tuple{InferenceState,LineNum}}(), # backedges
  Vector{InferenceState}(), # callers_in_cycle
  #=parent=#nothing,
@@ -238,9 +229,8 @@ function add_mt_backedge!(mt::Core.MethodTable, @nospecialize(typ), caller::Infe
  nothing
 end
 
-function is_specializable_vararg_slot(@nospecialize(arg), sv::InferenceState)
- return (isa(arg, Slot) && slot_id(arg) == sv.nargs &&
- isa(sv.vararg_type_container, DataType))
+function is_specializable_vararg_slot(@nospecialize(arg), nargs, vargs)
+ return (isa(arg, Slot) && slot_id(arg) == nargs && !isempty(vargs))
 end
 
 function print_callstack(sv::InferenceState)

base/compiler/optimize.jl

@@ -6,7 +6,7 @@
 
 mutable struct OptimizationState
  linfo::MethodInstance
- vararg_type_container #::Type
+ result_vargs::Vector{Any}
  backedges::Vector{Any}
  src::CodeInfo
  mod::Module
@@ -23,7 +23,7 @@ mutable struct OptimizationState
  end
  src = frame.src
  next_label = max(label_counter(src.code), length(src.code)) + 10
- return new(frame.linfo, frame.vararg_type_container,
+ return new(frame.linfo, frame.result.vargs,
  s_edges::Vector{Any},
  src, frame.mod, frame.nargs,
  next_label, frame.min_valid, frame.max_valid,
@@ -53,8 +53,8 @@ mutable struct OptimizationState
  nargs = 0
  end
  next_label = max(label_counter(src.code), length(src.code)) + 10
- vararg_type_container = nothing # if you want something more accurate, set it yourself :P
- return new(linfo, vararg_type_container,
+ result_vargs = Any[] # if you want something more accurate, set it yourself :P
+ return new(linfo, result_vargs,
  s_edges::Vector{Any},
  src, inmodule, nargs,
  next_label,
@@ -100,11 +100,6 @@ function add_backedge!(li::MethodInstance, caller::OptimizationState)
  nothing
 end
 
-function is_specializable_vararg_slot(@nospecialize(arg), sv::OptimizationState)
- return (isa(arg, Slot) && slot_id(arg) == sv.nargs &&
- isa(sv.vararg_type_container, DataType))
-end
-
 ###########
 # structs #
 ###########
@@ -1333,22 +1328,30 @@ function inlineable(@nospecialize(f), @nospecialize(ft), e::Expr, atypes::Vector
  if invoke_api(linfo) == 2
  # in this case function can be inlined to a constant
  add_backedge!(linfo, sv)
+ # XXX: @vtjnash thinks this should be `argexprs0`, but doing so exposes a
+ # downstream optimizer problem that breaks tests, so we're going to avoid
+ # changing it for now. ref
+ # https://github.com/JuliaLang/julia/pull/26826#issuecomment-386381103
  return inline_as_constant(linfo.inferred_const, argexprs, sv, invoke_data)
  end
 
- # see if the method has a InferenceResult in the current cache
+ # See if the method has a InferenceResult in the current cache
  # or an existing inferred code info store in `.inferred`
+ #
+ # Above, we may have rewritten trailing varargs in `atypes` to a tuple type. However,
+ # inference populates the cache with the pre-rewrite version (`atypes0`), so here, we
+ # check against that instead.
  haveconst = false
- for i in 1:length(atypes)
- a = atypes[i]
+ for i in 1:length(atypes0)
+ a = atypes0[i]
  if isa(a, Const) && !isdefined(typeof(a.val), :instance) && !(isa(a.val, Type) && issingletontype(a.val))
  # have new information from argtypes that wasn't available from the signature
  haveconst = true
  break
  end
  end
  if haveconst
- inf_result = cache_lookup(linfo, atypes, sv.params.cache) # Union{Nothing, InferenceResult}
+ inf_result = cache_lookup(linfo, atypes0, sv.params.cache) # Union{Nothing, InferenceResult}
  else
  inf_result = nothing
  end
@@ -2003,8 +2006,8 @@ function inline_call(e::Expr, sv::OptimizationState, stmts::Vector{Any}, boundsc
  tmpv = newvar!(sv, t)
  push!(newstmts, Expr(:(=), tmpv, aarg))
  end
- if is_specializable_vararg_slot(aarg, sv)
- tp = sv.vararg_type_container.parameters
+ if is_specializable_vararg_slot(aarg, sv.nargs, sv.result_vargs)
+ tp = sv.result_vargs
  else
  tp = t.parameters
  end

base/compiler/tfuncs.jl

@@ -223,7 +223,7 @@ add_tfunc(===, 2, 2,
  end
  return Bool
  end, 1)
-function isdefined_tfunc(args...)
+function isdefined_tfunc(@nospecialize(args...))
  arg1 = args[1]
  if isa(arg1, Const)
  a1 = typeof(arg1.val)

base/compiler/typelimits.jl

@@ -50,15 +50,12 @@ function is_derived_type(@nospecialize(t), @nospecialize(c), mindepth::Int)
  if t === c
  return mindepth == 0
  end
- if isa(c, TypeVar)
- # see if it is replacing a TypeVar upper bound with something simpler
- return is_derived_type(t, c.ub, mindepth)
- elseif isa(c, Union)
+ if isa(c, Union)
  # see if it is one of the elements of the union
  return is_derived_type(t, c.a, mindepth + 1) || is_derived_type(t, c.b, mindepth + 1)
  elseif isa(c, UnionAll)
  # see if it is derived from the body
- return is_derived_type(t, c.body, mindepth)
+ return is_derived_type(t, c.var.ub, mindepth) || is_derived_type(t, c.body, mindepth + 1)
  elseif isa(c, DataType)
  if isa(t, DataType)
  # see if it is one of the supertypes of a parameter
@@ -96,7 +93,8 @@ function is_derived_type_from_any(@nospecialize(t), sources::SimpleVector, minde
  return false
 end
 
-# type vs. comparison or which was derived from source
+# The goal of this function is to return a type of greater "size" and less "complexity" than
+# both `t` or `c` over the lattice defined by `sources`, `depth`, and `allowed_tuplelen`.
 function _limit_type_size(@nospecialize(t), @nospecialize(c), sources::SimpleVector, depth::Int, allowed_tuplelen::Int)
  if t === c
  return t # quick egal test
@@ -140,9 +138,9 @@ function _limit_type_size(@nospecialize(t), @nospecialize(c), sources::SimpleVec
  lb = Bottom
  end
  v2 = TypeVar(tv.name, lb, ub)
- return UnionAll(v2, _limit_type_size(t{v2}, c{v2}, sources, depth + 1, allowed_tuplelen))
+ return UnionAll(v2, _limit_type_size(t{v2}, c{v2}, sources, depth, allowed_tuplelen))
  end
- tbody = _limit_type_size(t.body, c, sources, depth + 1, allowed_tuplelen)
+ tbody = _limit_type_size(t.body, c, sources, depth, allowed_tuplelen)
  tbody === t.body && return t
  return UnionAll(t.var, tbody)
  elseif isa(c, UnionAll)

base/compiler/typeutils.jl

@@ -99,20 +99,6 @@ function tuple_tail_elem(@nospecialize(init), ct)
  return Vararg{widenconst(foldl((a, b) -> tmerge(a, tvar_extent(unwrapva(b))), init, ct))}
 end
 
-# t[n:end]
-function tupleparam_tail(t::SimpleVector, n)
- lt = length(t)
- if n > lt
- va = t[lt]
- if isvarargtype(va)
- # assumes that we should never see Vararg{T, x}, where x is a constant (should be guaranteed by construction)
- return Tuple{va}
- end
- return Tuple{}
- end
- return Tuple{t[n:lt]...}
-end
-
 # take a Tuple where one or more parameters are Unions
 # and return an array such that those Unions are removed
 # and `Union{return...} == ty`