Merge pull request #30878 from JuliaLang/jb/sptypes

some improvements to static parameter handling in inference

base/compiler/abstractinterpretation.jl

@@ -863,21 +863,6 @@ function abstract_eval_cfunction(e::Expr, vtypes::VarTable, sv::InferenceState)
  nothing
 end
 
-# convert an inferred static parameter value to the inferred type of a static_parameter expression
-function sparam_type(@nospecialize(val))
- if isa(val, TypeVar)
- if Any <: val.ub
- # static param bound to typevar
- # if the tvar is not known to refer to anything more specific than Any,
- # the static param might actually be an integer, symbol, etc.
- return Any
- else
- return UnionAll(val, Type{val})
- end
- end
- return AbstractEvalConstant(val)
-end
-
 function abstract_eval(@nospecialize(e), vtypes::VarTable, sv::InferenceState)
  if isa(e, QuoteNode)
  return AbstractEvalConstant((e::QuoteNode).value)
@@ -940,8 +925,8 @@ function abstract_eval(@nospecialize(e), vtypes::VarTable, sv::InferenceState)
  elseif e.head === :static_parameter
  n = e.args[1]
  t = Any
- if 1 <= n <= length(sv.sp)
- t = sparam_type(sv.sp[n])
+ if 1 <= n <= length(sv.sptypes)
+ t = sv.sptypes[n]
  end
  elseif e.head === :method
  t = (length(e.args) == 1) ? Any : Nothing
@@ -975,9 +960,9 @@ function abstract_eval(@nospecialize(e), vtypes::VarTable, sv::InferenceState)
  end
  elseif isa(sym, Expr) && sym.head === :static_parameter
  n = sym.args[1]
- if 1 <= n <= length(sv.sp)
- val = sv.sp[n]
- if !isa(val, TypeVar)
+ if 1 <= n <= length(sv.sptypes)
+ spty = sv.sptypes[n]
+ if isa(spty, Const)
  t = Const(true)
  end
  end

base/compiler/inferencestate.jl

@@ -5,8 +5,8 @@ const LineNum = Int
 mutable struct InferenceState
  params::Params # describes how to compute the result
  result::InferenceResult # remember where to put the result
- linfo::MethodInstance # used here for the tuple (specTypes, env, Method) and world-age validity
- sp::SimpleVector  # static parameters
+ linfo::MethodInstance  # used here for the tuple (specTypes, env, Method) and world-age validity
+ sptypes::Vector{Any} # types of static parameter
  slottypes::Vector{Any}
  mod::Module
  currpc::LineNum
@@ -48,7 +48,7 @@ mutable struct InferenceState
  code = src.code::Array{Any,1}
  toplevel = !isa(linfo.def, Method)
 
- sp = spvals_from_meth_instance(linfo::MethodInstance)
+ sp = sptypes_from_meth_instance(linfo::MethodInstance)
 
  nssavalues = src.ssavaluetypes::Int
  src.ssavaluetypes = Any[ NOT_FOUND for i = 1:nssavalues ]
@@ -120,7 +120,7 @@ function InferenceState(result::InferenceResult, cached::Bool, params::Params)
  return InferenceState(result, src, cached, params)
 end
 
-function spvals_from_meth_instance(linfo::MethodInstance)
+function sptypes_from_meth_instance(linfo::MethodInstance)
  toplevel = !isa(linfo.def, Method)
  if !toplevel && isempty(linfo.sparam_vals) && !isempty(linfo.def.sparam_syms)
  # linfo is unspecialized
@@ -130,35 +130,54 @@ function spvals_from_meth_instance(linfo::MethodInstance)
  push!(sp, sig.var)
  sig = sig.body
  end
- sp = svec(sp...)
  else
- sp = linfo.sparam_vals
- if _any(t->isa(t,TypeVar), sp)
- sp = collect(Any, sp)
- end
+ sp = collect(Any, linfo.sparam_vals)
  end
- if !isa(sp, SimpleVector)
- for i = 1:length(sp)
- v = sp[i]
- if v isa TypeVar
- ub = v.ub
- while ub isa TypeVar
- ub = ub.ub
- end
- if has_free_typevars(ub)
- ub = Any
+ for i = 1:length(sp)
+ v = sp[i]
+ if v isa TypeVar
+ ub = v.ub
+ while ub isa TypeVar
+ ub = ub.ub
+ end
+ if has_free_typevars(ub)
+ ub = Any
+ end
+ lb = v.lb
+ while lb isa TypeVar
+ lb = lb.lb
+ end
+ if has_free_typevars(lb)
+ lb = Bottom
+ end
+ if Any <: ub && lb <: Bottom
+ ty = Any
+ # if this parameter came from arg::Type{T}, we know that T::Type
+ sig = linfo.def.sig
+ temp = sig
+ for j = 1:i-1
+ temp = temp.body
  end
- lb = v.lb
- while lb isa TypeVar
- lb = lb.lb
+ Pi = temp.var
+ while temp isa UnionAll
+ temp = temp.body
  end
- if has_free_typevars(lb)
- lb = Bottom
+ sigtypes = temp.parameters
+ for j = 1:length(sigtypes)
+ tj = sigtypes[j]
+ if isType(tj) && tj.parameters[1] === Pi
+ ty = Type
+ break
+ end
  end
- sp[i] = TypeVar(v.name, lb, ub)
+ else
+ tv = TypeVar(v.name, lb, ub)
+ ty = UnionAll(tv, Type{tv})
  end
+ else
+ ty = Const(v)
  end
- sp = svec(sp...)
+ sp[i] = ty
  end
  return sp
 end

base/compiler/optimize.jl

@@ -13,7 +13,7 @@ mutable struct OptimizationState
  min_valid::UInt
  max_valid::UInt
  params::Params
- sp::SimpleVector # static parameters
+ sptypes::Vector{Any} # static parameters
  slottypes::Vector{Any}
  const_api::Bool
  function OptimizationState(frame::InferenceState)
@@ -27,7 +27,7 @@ mutable struct OptimizationState
  s_edges::Vector{Any},
  src, frame.mod, frame.nargs,
  frame.min_valid, frame.max_valid,
- frame.params, frame.sp, frame.slottypes, false)
+ frame.params, frame.sptypes, frame.slottypes, false)
  end
  function OptimizationState(linfo::MethodInstance, src::CodeInfo,
  params::Params)
@@ -54,7 +54,7 @@ mutable struct OptimizationState
  s_edges::Vector{Any},
  src, inmodule, nargs,
  min_world(linfo), max_world(linfo),
- params, spvals_from_meth_instance(linfo), slottypes, false)
+ params, sptypes_from_meth_instance(linfo), slottypes, false)
  end
 end
 
@@ -135,7 +135,7 @@ function isinlineable(m::Method, me::OptimizationState, bonus::Int=0)
  end
  end
  if !inlineable
- inlineable = inline_worthy(me.src.code, me.src, me.sp, me.slottypes, me.params, cost_threshold + bonus)
+ inlineable = inline_worthy(me.src.code, me.src, me.sptypes, me.slottypes, me.params, cost_threshold + bonus)
  end
  return inlineable
 end
@@ -148,7 +148,7 @@ function stmt_affects_purity(@nospecialize(stmt), ir)
  return false
  end
  if isa(stmt, GotoIfNot)
- t = argextype(stmt.cond, ir, ir.spvals)
+ t = argextype(stmt.cond, ir, ir.sptypes)
  return !(t ⊑ Bool)
  end
  if isa(stmt, Expr)
@@ -175,7 +175,7 @@ function optimize(opt::OptimizationState, @nospecialize(result))
  proven_pure = true
  for i in 1:length(ir.stmts)
  stmt = ir.stmts[i]
- if stmt_affects_purity(stmt, ir) && !stmt_effect_free(stmt, ir.types[i], ir, ir.spvals)
+ if stmt_affects_purity(stmt, ir) && !stmt_effect_free(stmt, ir.types[i], ir, ir.sptypes)
  proven_pure = false
  break
  end
@@ -268,19 +268,19 @@ plus_saturate(x::Int, y::Int) = max(x, y, x+y)
 # known return type
 isknowntype(@nospecialize T) = (T == Union{}) || isconcretetype(T)
 
-function statement_cost(ex::Expr, line::Int, src::CodeInfo, spvals::SimpleVector, slottypes::Vector{Any}, params::Params)
+function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}, slottypes::Vector{Any}, params::Params)
  head = ex.head
  if is_meta_expr_head(head)
  return 0
  elseif head === :call
  farg = ex.args[1]
- ftyp = argextype(farg, src, spvals, slottypes)
+ ftyp = argextype(farg, src, sptypes, slottypes)
  if ftyp === IntrinsicFunction && farg isa SSAValue
  # if this comes from code that was already inlined into another function,
  # Consts have been widened. try to recover in simple cases.
  farg = src.code[farg.id]
  if isa(farg, GlobalRef) || isa(farg, QuoteNode) || isa(farg, IntrinsicFunction) || isexpr(farg, :static_parameter)
- ftyp = argextype(farg, src, spvals, slottypes)
+ ftyp = argextype(farg, src, sptypes, slottypes)
  end
  end
  f = singleton_type(ftyp)
@@ -302,7 +302,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, spvals::SimpleVector
  # return plus_saturate(argcost, isknowntype(extyp) ? 1 : params.inline_nonleaf_penalty)
  return 0
  elseif f === Main.Core.arrayref && length(ex.args) >= 3
- atyp = argextype(ex.args[3], src, spvals, slottypes)
+ atyp = argextype(ex.args[3], src, sptypes, slottypes)
  return isknowntype(atyp) ? 4 : params.inline_nonleaf_penalty
  end
  fidx = find_tfunc(f)
@@ -325,7 +325,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, spvals::SimpleVector
  elseif head === :return
  a = ex.args[1]
  if a isa Expr
- return statement_cost(a, -1, src, spvals, slottypes, params)
+ return statement_cost(a, -1, src, sptypes, slottypes, params)
  end
  return 0
  elseif head === :(=)
@@ -336,7 +336,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, spvals::SimpleVector
  end
  a = ex.args[2]
  if a isa Expr
- cost = plus_saturate(cost, statement_cost(a, -1, src, spvals, slottypes, params))
+ cost = plus_saturate(cost, statement_cost(a, -1, src, sptypes, slottypes, params))
  end
  return cost
  elseif head === :copyast
@@ -357,13 +357,13 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, spvals::SimpleVector
  return 0
 end
 
-function inline_worthy(body::Array{Any,1}, src::CodeInfo, spvals::SimpleVector, slottypes::Vector{Any},
+function inline_worthy(body::Array{Any,1}, src::CodeInfo, sptypes::Vector{Any}, slottypes::Vector{Any},
  params::Params, cost_threshold::Integer=params.inline_cost_threshold)
  bodycost::Int = 0
  for line = 1:length(body)
  stmt = body[line]
  if stmt isa Expr
- thiscost = statement_cost(stmt, line, src, spvals, slottypes, params)::Int
+ thiscost = statement_cost(stmt, line, src, sptypes, slottypes, params)::Int
  elseif stmt isa GotoNode
  # loops are generally always expensive
  # but assume that forward jumps are already counted for from
@@ -378,11 +378,11 @@ function inline_worthy(body::Array{Any,1}, src::CodeInfo, spvals::SimpleVector,
  return true
 end
 
-function is_known_call(e::Expr, @nospecialize(func), src, spvals::SimpleVector, slottypes::Vector{Any} = empty_slottypes)
+function is_known_call(e::Expr, @nospecialize(func), src, sptypes::Vector{Any}, slottypes::Vector{Any} = empty_slottypes)
  if e.head !== :call
  return false
  end
- f = argextype(e.args[1], src, spvals, slottypes)
+ f = argextype(e.args[1], src, sptypes, slottypes)
  return isa(f, Const) && f.val === func
 end
 

base/compiler/ssair/driver.jl

@@ -104,9 +104,9 @@ function just_construct_ssa(ci::CodeInfo, code::Vector{Any}, nargs::Int, sv::Opt
  @timeit "domtree 1" domtree = construct_domtree(cfg)
  ir = let code = Any[nothing for _ = 1:length(code)]
  argtypes = sv.slottypes[1:(nargs+1)]
- IRCode(code, Any[], ci.codelocs, flags, cfg, collect(LineInfoNode, ci.linetable), argtypes, meta, sv.sp)
+ IRCode(code, Any[], ci.codelocs, flags, cfg, collect(LineInfoNode, ci.linetable), argtypes, meta, sv.sptypes)
  end
- @timeit "construct_ssa" ir = construct_ssa!(ci, code, ir, domtree, defuse_insts, nargs, sv.sp, sv.slottypes)
+ @timeit "construct_ssa" ir = construct_ssa!(ci, code, ir, domtree, defuse_insts, nargs, sv.sptypes, sv.slottypes)
  return ir
 end
 

base/compiler/ssair/inlining.jl

@@ -787,7 +787,7 @@ function assemble_inline_todo!(ir::IRCode, linetable::Vector{LineInfoNode}, sv::
  isempty(eargs) && continue
  arg1 = eargs[1]
 
- ft = argextype(arg1, ir, sv.sp)
+ ft = argextype(arg1, ir, sv.sptypes)
  has_free_typevars(ft) && continue
  f = singleton_type(ft)
  f === Core.Intrinsics.llvmcall && continue
@@ -797,7 +797,7 @@ function assemble_inline_todo!(ir::IRCode, linetable::Vector{LineInfoNode}, sv::
  atypes[1] = ft
  ok = true
  for i = 2:length(stmt.args)
- a = argextype(stmt.args[i], ir, sv.sp)
+ a = argextype(stmt.args[i], ir, sv.sptypes)
  (a === Bottom || isvarargtype(a)) && (ok = false; break)
  atypes[i] = a
  end

base/compiler/ssair/ir.jl

@@ -213,20 +213,20 @@ struct IRCode
  lines::Vector{Int32}
  flags::Vector{UInt8}
  argtypes::Vector{Any}
- spvals::SimpleVector
+ sptypes::Vector{Any}
  linetable::Vector{LineInfoNode}
  cfg::CFG
  new_nodes::Vector{NewNode}
  meta::Vector{Any}
 
  function IRCode(stmts::Vector{Any}, types::Vector{Any}, lines::Vector{Int32}, flags::Vector{UInt8},
  cfg::CFG, linetable::Vector{LineInfoNode}, argtypes::Vector{Any}, meta::Vector{Any},
- spvals::SimpleVector)
- return new(stmts, types, lines, flags, argtypes, spvals, linetable, cfg, NewNode[], meta)
+ sptypes::Vector{Any})
+ return new(stmts, types, lines, flags, argtypes, sptypes, linetable, cfg, NewNode[], meta)
  end
  function IRCode(ir::IRCode, stmts::Vector{Any}, types::Vector{Any}, lines::Vector{Int32}, flags::Vector{UInt8},
  cfg::CFG, new_nodes::Vector{NewNode})
- return new(stmts, types, lines, flags, ir.argtypes, ir.spvals, ir.linetable, cfg, new_nodes, ir.meta)
+ return new(stmts, types, lines, flags, ir.argtypes, ir.sptypes, ir.linetable, cfg, new_nodes, ir.meta)
  end
 end
 copy(code::IRCode) = IRCode(code, copy(code.stmts), copy(code.types),
@@ -1143,7 +1143,7 @@ function maybe_erase_unused!(extra_worklist, compact, idx, callback = x->nothing
  if compact_exprtype(compact, SSAValue(idx)) === Bottom
  effect_free = false
  else
- effect_free = stmt_effect_free(stmt, compact.result_types[idx], compact, compact.ir.spvals)
+ effect_free = stmt_effect_free(stmt, compact.result_types[idx], compact, compact.ir.sptypes)
  end
  if effect_free
  for ops in userefs(stmt)

base/compiler/ssair/legacy.jl

@@ -1,18 +1,18 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
-inflate_ir(ci::CodeInfo) = inflate_ir(ci, Core.svec(), Any[ Any for i = 1:length(ci.slotnames) ])
+inflate_ir(ci::CodeInfo) = inflate_ir(ci, Any[], Any[ Any for i = 1:length(ci.slotnames) ])
 
 function inflate_ir(ci::CodeInfo, linfo::MethodInstance)
- spvals = spvals_from_meth_instance(linfo)
+ sptypes = sptypes_from_meth_instance(linfo)
  if ci.inferred
  argtypes, _ = matching_cache_argtypes(linfo, nothing)
  else
  argtypes = Any[ Any for i = 1:length(ci.slotnames) ]
  end
- return inflate_ir(ci, spvals, argtypes)
+ return inflate_ir(ci, sptypes, argtypes)
 end
 
-function inflate_ir(ci::CodeInfo, spvals::SimpleVector, argtypes::Vector{Any})
+function inflate_ir(ci::CodeInfo, sptypes::Vector{Any}, argtypes::Vector{Any})
  code = copy_exprargs(ci.code)
  for i = 1:length(code)
  if isa(code[i], Expr)
@@ -46,7 +46,7 @@ function inflate_ir(ci::CodeInfo, spvals::SimpleVector, argtypes::Vector{Any})
  end
  ssavaluetypes = ci.ssavaluetypes isa Vector{Any} ? copy(ci.ssavaluetypes) : Any[ Any for i = 1:(ci.ssavaluetypes::Int) ]
  ir = IRCode(code, ssavaluetypes, copy(ci.codelocs), copy(ci.ssaflags), cfg, collect(LineInfoNode, ci.linetable),
- argtypes, Any[], spvals)
+ argtypes, Any[], sptypes)
  return ir
 end