inference×effects: improve the const-prop' heuristics

This commit improves the heuristics to judge const-prop' profitability
with new effect property `:const_prop_profitable_args`. This is supposed
to supplement our primary const-prop' heuristic based on inlining cost
and is supposed to be a general fix for type stabilities issues discussed
at e.g. #45952 and #46430 (and eliminating the need for manual
`@constprop :aggressive` clutters in such situations).

The new effect property `:const_prop_profitable_args` tracks call
arguments that can be considered to shape up generated code if
their constant information is available. Currently this commit
exploits the following const-prop' profitabilities:
- `Val(x)`-profitability: as `Val` generally encodes constant information
  into the type domain, it is generally profitable to constant prop' `x`
  if the constructed `Val(x)` is used later (e.g. for dispatch).
  This basically tries to exploit const-prop' profitability in the
  following kind of case:
  ```julia
  kernel(::Val{1}, args...) = ...
  kernel(::Val{2}, args...) = ...

  function profitable1(x::Int, args...)
      kernel(Val(x), args...)
  end
  ```
  This allows the compiler to perform const-prop' for case like #45952
  even if the primary heuristics based on inlining cost gets confused.
- branching-profitability: constant branch condition is generally very
  profitable as it can shape up generated code as well as narrow down
  the return type inference by cutting off the dead branch.
  ```julia
  function profitable2(raise::Bool, args...)
      v = op(args...)
      if v === nothing && raise
          return nothing
      end
      return v
  end
  ```

Currently this commit passes all the test cases and also actually
improves target type stabilities, but doesn't work very ideally as it
seems to be a bit too aggressive (this commit right now strictly
increases the chances of const-propagation). I'd like to further tweak
this heuristic to keep the latency in general cases.

Author: aviatesk

base/compiler/abstractinterpretation.jl

@@ -210,7 +210,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
         all_effects = Effects(all_effects; nothrow=false)
     end
 
-    rettype = from_interprocedural!(𝕃ₚ, rettype, sv, arginfo, conditionals)
+    (; rt, effects) = from_interprocedural!(𝕃ₚ, rettype, all_effects, sv, arginfo, conditionals)
 
     # Also considering inferring the compilation signature for this method, so
     # it is available to the compiler in case it ends up needing it.
@@ -219,32 +219,32 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
         method = match.method
         sig = match.spec_types
         mi = specialize_method(match; preexisting=true)
-        if mi !== nothing && !const_prop_methodinstance_heuristic(interp, mi, arginfo, sv)
+        if mi !== nothing && !const_prop_methodinstance_heuristic(interp, mi, arginfo, Effects(), sv)
             csig = get_compileable_sig(method, sig, match.sparams)
             if csig !== nothing && csig !== sig
                 abstract_call_method(interp, method, csig, match.sparams, multiple_matches, StmtInfo(false), sv)
             end
         end
     end
 
-    if call_result_unused(si) && !(rettype === Bottom)
+    if call_result_unused(si) && !(rt === Bottom)
         add_remark!(interp, sv, "Call result type was widened because the return value is unused")
         # We're mainly only here because the optimizer might want this code,
         # but we ourselves locally don't typically care about it locally
         # (beyond checking if it always throws).
         # So avoid adding an edge, since we don't want to bother attempting
         # to improve our result even if it does change (to always throw),
         # and avoid keeping track of a more complex result type.
-        rettype = Any
+        rt = Any
     end
-    add_call_backedges!(interp, rettype, all_effects, edges, matches, atype, sv)
+    add_call_backedges!(interp, rt, effects, edges, matches, atype, sv)
     if !isempty(sv.pclimitations) # remove self, if present
         delete!(sv.pclimitations, sv)
         for caller in sv.callers_in_cycle
             delete!(sv.pclimitations, caller)
         end
     end
-    return CallMeta(rettype, all_effects, info)
+    return CallMeta(rt, effects, info)
 end
 
 struct FailedMethodMatch
@@ -348,15 +348,24 @@ function find_matching_methods(𝕃::AbstractLattice,
     end
 end
 
+struct InterproceduralResult
+    rt
+    effects::Effects
+    InterproceduralResult(@nospecialize(rt), effects::Effects) = new(rt, effects)
+end
+
 """
-    from_interprocedural!(𝕃ₚ::AbstractLattice, rt, sv::InferenceState, arginfo::ArgInfo, maybecondinfo) -> newrt
+    from_interprocedural!(𝕃ₚ::AbstractLattice, rt, effects::Effects,
+        sv::InferenceState, arginfo::ArgInfo, maybecondinfo) -> InterproceduralResult
 
-Converts inter-procedural return type `rt` into a local lattice element `newrt`,
-that is appropriate in the context of current local analysis frame `sv`, especially:
+Converts extended lattice element `rt` and `effects::Effects` that represent inferred
+return type and method call effects into new lattice ement and `Effects` that are
+appropriate in the context of current local analysis frame `sv`, especially:
 - unwraps `rt::LimitedAccuracy` and collects its limitations into the current frame `sv`
 - converts boolean `rt` to new boolean `newrt` in a way `newrt` can propagate extra conditional
   refinement information, e.g. translating `rt::InterConditional` into `newrt::Conditional`
   that holds a type constraint information about a variable in `sv`
+- recomputes `effects.const_prop_profitable_args` so that they are imposed on call arguments of `sv`
 
 This function _should_ be used wherever we propagate results returned from
 `abstract_call_method` or `abstract_call_method_with_const_args`.
@@ -368,7 +377,8 @@ In such cases `maybecondinfo` should be either of:
 When we deal with multiple `MethodMatch`es, it's better to precompute `maybecondinfo` by
 `tmerge`ing argument signature type of each method call.
 """
-function from_interprocedural!(𝕃ₚ::AbstractLattice, @nospecialize(rt), sv::InferenceState, arginfo::ArgInfo, @nospecialize(maybecondinfo))
+function from_interprocedural!(𝕃ₚ::AbstractLattice, @nospecialize(rt), effects::Effects,
+    sv::InferenceState, arginfo::ArgInfo, @nospecialize(maybecondinfo))
     rt = collect_limitations!(rt, sv)
     if isa(rt, InterMustAlias)
         rt = from_intermustalias(rt, arginfo)
@@ -380,7 +390,23 @@ function from_interprocedural!(𝕃ₚ::AbstractLattice, @nospecialize(rt), sv::
         end
     end
     @assert !(rt isa InterConditional || rt isa InterMustAlias) "invalid lattice element returned from inter-procedural context"
-    return rt
+    if effects.const_prop_profitable_args !== NO_PROFITABLE_ARGS
+        argsbits = 0x00
+        fargs = arginfo.fargs
+        if fargs !== nothing
+            for i = 1:length(fargs)
+                if is_const_prop_profitable_arg(effects, i)
+                    arg = fargs[i]
+                    if is_call_argument(arg, sv) && 1 ≤ slot_id(arg) ≤ 8
+                        argsbits |= 0x01 << (slot_id(arg)-1)
+                    end
+                end
+            end
+        end
+        const_prop_profitable_args = ConstPropProfitableArgs(argsbits)
+        effects = Effects(effects; const_prop_profitable_args)
+    end
+    return InterproceduralResult(rt, effects)
 end
 
 function collect_limitations!(@nospecialize(typ), sv::InferenceState)
@@ -993,9 +1019,8 @@ function abstract_call_method_with_const_args(interp::AbstractInterpreter,
         end
     end
     # try constant prop'
-    inf_cache = get_inference_cache(interp)
     𝕃ᵢ = typeinf_lattice(interp)
-    inf_result = cache_lookup(𝕃ᵢ, mi, arginfo.argtypes, inf_cache)
+    inf_result = cache_lookup(𝕃ᵢ, mi, arginfo.argtypes, get_inference_cache(interp))
     if inf_result === nothing
         # if there might be a cycle, check to make sure we don't end up
         # calling ourselves here.
@@ -1062,7 +1087,7 @@ function maybe_get_const_prop_profitable(interp::AbstractInterpreter,
         return nothing
     end
     mi = mi::MethodInstance
-    if !force && !const_prop_methodinstance_heuristic(interp, mi, arginfo, sv)
+    if !force && !const_prop_methodinstance_heuristic(interp, mi, arginfo, result.effects, sv)
         add_remark!(interp, sv, "[constprop] Disabled by method instance heuristic")
         return nothing
     end
@@ -1214,7 +1239,7 @@ end
 # where we would spend a lot of time, but are probably unlikely to get an improved
 # result anyway.
 function const_prop_methodinstance_heuristic(interp::AbstractInterpreter,
-    mi::MethodInstance, arginfo::ArgInfo, sv::InferenceState)
+    mi::MethodInstance, arginfo::ArgInfo, effects::Effects, sv::InferenceState)
     method = mi.def::Method
     if method.is_for_opaque_closure
         # Not inlining an opaque closure can be very expensive, so be generous
@@ -1239,6 +1264,8 @@ function const_prop_methodinstance_heuristic(interp::AbstractInterpreter,
     elseif is_stmt_noinline(flag)
         # this call won't be inlined, thus this constant-prop' will most likely be unfruitful
         return false
+    elseif any_const_prop_profitable_args(effects, arginfo.argtypes)
+        return true
     else
         # Peek at the inferred result for the method to determine if the optimizer
         # was able to cut it down to something simple (inlineable in particular).
@@ -1256,6 +1283,21 @@ function const_prop_methodinstance_heuristic(interp::AbstractInterpreter,
     return false # the cache isn't inlineable, so this constant-prop' will most likely be unfruitful
 end
 
+# check if constant information is available on any call argument that has been analyzed as
+# const-prop' profitable
+function any_const_prop_profitable_args(effects::Effects, argtypes::Vector{Any})
+    if effects.const_prop_profitable_args === NO_PROFITABLE_ARGS
+        return false
+    end
+    for i in 1:length(argtypes)
+        ai = widenconditional(argtypes[i])
+        if isa(ai, Const) && is_const_prop_profitable_arg(effects, i)
+            return true
+        end
+    end
+    return false
+end
+
 # This is only for use with `Conditional`.
 # In general, usage of this is wrong.
 ssa_def_slot(@nospecialize(arg), sv::IRCode) = nothing
@@ -1901,11 +1943,10 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn
             (; rt, effects, const_result, edge) = const_call_result
         end
     end
-    rt = from_interprocedural!(𝕃ₚ, rt, sv, arginfo, sig)
     effects = Effects(effects; nonoverlayed=!overlayed)
-    info = InvokeCallInfo(match, const_result)
+    (; rt, effects) = from_interprocedural!(𝕃ₚ, rt, effects, sv, arginfo, sig)
     edge !== nothing && add_invoke_backedge!(sv, lookupsig, edge)
-    return CallMeta(rt, effects, info)
+    return CallMeta(rt, effects, InvokeCallInfo(match, const_result))
 end
 
 function invoke_rewrite(xs::Vector{Any})
@@ -2015,6 +2056,20 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
         return CallMeta(typename_static(argtypes[2]), EFFECTS_TOTAL, MethodResultPure())
     elseif f === Core._hasmethod
         return _hasmethod_tfunc(interp, argtypes, sv)
+    elseif la == 2 && istoptype(f, :Val)
+        # `Val` generally encodes constant information into the type domain, so there is
+        # generally a high profitability for constant propagation if the argument of the
+        # `Val` constructor is a call argument
+        fargs = arginfo.fargs
+        if fargs !== nothing
+            arg = arginfo.fargs[2]
+            if is_call_argument(arg, sv) && !isempty(sv.ssavalue_uses[sv.currpc])
+                if 1 ≤ slot_id(arg) ≤ 8
+                    const_prop_profitable_args = ConstPropProfitableArgs(0x01 << (slot_id(arg)-1))
+                    merge_effects!(interp, sv, Effects(EFFECTS_TOTAL; const_prop_profitable_args))
+                end
+            end
+        end
     end
     atype = argtypes_to_type(argtypes)
     return abstract_call_gf_by_type(interp, f, arginfo, si, atype, sv, max_methods)
@@ -2048,7 +2103,7 @@ function abstract_call_opaque_closure(interp::AbstractInterpreter,
             effects = Effects(effects; nothrow=false)
         end
     end
-    rt = from_interprocedural!(𝕃ₚ, rt, sv, arginfo, match.spec_types)
+    (; rt, effects) = from_interprocedural!(𝕃ₚ, rt, effects, sv, arginfo, match.spec_types)
     info = OpaqueClosureCallInfo(match, const_result)
     edge !== nothing && add_backedge!(sv, edge)
     return CallMeta(rt, effects, info)
@@ -2481,8 +2536,7 @@ function abstract_eval_foreigncall(interp::AbstractInterpreter, e::Expr, vtypes:
             override.terminates_globally ? true        : effects.terminates,
             override.notaskstate         ? true        : effects.notaskstate,
             override.inaccessiblememonly ? ALWAYS_TRUE : effects.inaccessiblememonly,
-            effects.nonoverlayed,
-            effects.noinbounds)
+            effects.nonoverlayed, effects.noinbounds, effects.const_prop_profitable_args)
     end
     return RTEffects(t, effects)
 end
@@ -2865,6 +2919,15 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
                     @goto branch
                 elseif isa(stmt, GotoIfNot)
                     condx = stmt.cond
+                    if is_call_argument(condx, frame)
+                        # if this condition object is a call argument, there will be a high
+                        # profitability for constant-propagating it, since it can shape up
+                        # the generated code by cutting off the dead branch entirely
+                        if 1 ≤ slot_id(condx) ≤ 8
+                            const_prop_profitable_args = ConstPropProfitableArgs(0x01 << (slot_id(condx)-1))
+                            merge_effects!(interp, frame, Effects(EFFECTS_TOTAL; const_prop_profitable_args))
+                        end
+                    end
                     condt = abstract_eval_value(interp, condx, currstate, frame)
                     if condt === Bottom
                         ssavaluetypes[currpc] = Bottom

base/compiler/effects.jl

@@ -1,3 +1,7 @@
+struct ConstPropProfitableArgs
+    argsbits::UInt8
+end
+
 """
     effects::Effects
 
@@ -64,6 +68,7 @@ struct Effects
     inaccessiblememonly::UInt8
     nonoverlayed::Bool
     noinbounds::Bool
+    const_prop_profitable_args::ConstPropProfitableArgs
     function Effects(
         consistent::UInt8,
         effect_free::UInt8,
@@ -72,7 +77,8 @@ struct Effects
         notaskstate::Bool,
         inaccessiblememonly::UInt8,
         nonoverlayed::Bool,
-        noinbounds::Bool)
+        noinbounds::Bool,
+        const_prop_profitable_args::ConstPropProfitableArgs = NO_PROFITABLE_ARGS)
         return new(
             consistent,
             effect_free,
@@ -81,7 +87,8 @@ struct Effects
             notaskstate,
             inaccessiblememonly,
             nonoverlayed,
-            noinbounds)
+            noinbounds,
+            const_prop_profitable_args)
     end
 end
 
@@ -98,10 +105,13 @@ const EFFECT_FREE_IF_INACCESSIBLEMEMONLY = 0x01 << 1
 # :inaccessiblememonly bits
 const INACCESSIBLEMEM_OR_ARGMEMONLY = 0x01 << 1
 
-const EFFECTS_TOTAL    = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  true,  true,  true,  ALWAYS_TRUE,  true,  true)
-const EFFECTS_THROWS   = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  false, true,  true,  ALWAYS_TRUE,  true,  true)
-const EFFECTS_UNKNOWN  = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, true,  false)  # unknown mostly, but it's not overlayed at least (e.g. it's not a call)
-const _EFFECTS_UNKNOWN = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, false, false) # unknown really
+# :const_prop_profitable_args bits
+const NO_PROFITABLE_ARGS = ConstPropProfitableArgs(0x00)
+
+const EFFECTS_TOTAL    = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  true,  true,  true,  ALWAYS_TRUE,  true,  true,  NO_PROFITABLE_ARGS)
+const EFFECTS_THROWS   = Effects(ALWAYS_TRUE,  ALWAYS_TRUE,  false, true,  true,  ALWAYS_TRUE,  true,  true,  NO_PROFITABLE_ARGS)
+const EFFECTS_UNKNOWN  = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, true,  false, NO_PROFITABLE_ARGS)  # unknown mostly, but it's not overlayed at least (e.g. it's not a call)
+const _EFFECTS_UNKNOWN = Effects(ALWAYS_FALSE, ALWAYS_FALSE, false, false, false, ALWAYS_FALSE, false, false, NO_PROFITABLE_ARGS) # unknown really
 
 function Effects(e::Effects = _EFFECTS_UNKNOWN;
     consistent::UInt8 = e.consistent,
@@ -111,7 +121,8 @@ function Effects(e::Effects = _EFFECTS_UNKNOWN;
     notaskstate::Bool = e.notaskstate,
     inaccessiblememonly::UInt8 = e.inaccessiblememonly,
     nonoverlayed::Bool = e.nonoverlayed,
-    noinbounds::Bool = e.noinbounds)
+    noinbounds::Bool = e.noinbounds,
+    const_prop_profitable_args::ConstPropProfitableArgs = e.const_prop_profitable_args)
     return Effects(
         consistent,
         effect_free,
@@ -120,7 +131,8 @@ function Effects(e::Effects = _EFFECTS_UNKNOWN;
         notaskstate,
         inaccessiblememonly,
         nonoverlayed,
-        noinbounds)
+        noinbounds,
+        const_prop_profitable_args)
 end
 
 function merge_effects(old::Effects, new::Effects)
@@ -132,7 +144,8 @@ function merge_effects(old::Effects, new::Effects)
         merge_effectbits(old.notaskstate, new.notaskstate),
         merge_effectbits(old.inaccessiblememonly, new.inaccessiblememonly),
         merge_effectbits(old.nonoverlayed, new.nonoverlayed),
-        merge_effectbits(old.noinbounds, new.noinbounds))
+        merge_effectbits(old.noinbounds, new.noinbounds),
+        merge_effectbits(old.const_prop_profitable_args, new.const_prop_profitable_args))
 end
 
 function merge_effectbits(old::UInt8, new::UInt8)
@@ -142,6 +155,7 @@ function merge_effectbits(old::UInt8, new::UInt8)
     return old | new
 end
 merge_effectbits(old::Bool, new::Bool) = old & new
+merge_effectbits(old::ConstPropProfitableArgs, new::ConstPropProfitableArgs) = ConstPropProfitableArgs(old.argsbits | new.argsbits)
 
 is_consistent(effects::Effects)          = effects.consistent === ALWAYS_TRUE
 is_effect_free(effects::Effects)         = effects.effect_free === ALWAYS_TRUE
@@ -177,15 +191,18 @@ is_effect_free_if_inaccessiblememonly(effects::Effects) = !iszero(effects.effect
 
 is_inaccessiblemem_or_argmemonly(effects::Effects) = effects.inaccessiblememonly === INACCESSIBLEMEM_OR_ARGMEMONLY
 
+is_const_prop_profitable_arg(effects::Effects, arg::Int) = !iszero(effects.const_prop_profitable_args.argsbits & (0x01 << (arg-1)))
+
 function encode_effects(e::Effects)
-    return ((e.consistent          % UInt32) << 0) |
-           ((e.effect_free         % UInt32) << 3) |
-           ((e.nothrow             % UInt32) << 5) |
-           ((e.terminates          % UInt32) << 6) |
-           ((e.notaskstate         % UInt32) << 7) |
-           ((e.inaccessiblememonly % UInt32) << 8) |
-           ((e.nonoverlayed        % UInt32) << 10)|
-           ((e.noinbounds          % UInt32) << 11)
+    return ((e.consistent          % UInt32) << 0)  |
+           ((e.effect_free         % UInt32) << 3)  |
+           ((e.nothrow             % UInt32) << 5)  |
+           ((e.terminates          % UInt32) << 6)  |
+           ((e.notaskstate         % UInt32) << 7)  |
+           ((e.inaccessiblememonly % UInt32) << 8)  |
+           ((e.nonoverlayed        % UInt32) << 10) |
+           ((e.noinbounds          % UInt32) << 11) |
+           ((e.const_prop_profitable_args.argsbits % UInt32) << 12)
 end
 
 function decode_effects(e::UInt32)
@@ -197,7 +214,8 @@ function decode_effects(e::UInt32)
         _Bool((e >> 7) & 0x01),
         UInt8((e >> 8) & 0x03),
         _Bool((e >> 10) & 0x01),
-        _Bool((e >> 11) & 0x01))
+        _Bool((e >> 11) & 0x01),
+        ConstPropProfitableArgs(UInt8((e >> 12) & 0x7f)))
 end
 
 struct EffectsOverride

base/compiler/inferencestate.jl

@@ -620,3 +620,5 @@ function narguments(sv::InferenceState, include_va::Bool=true)
     end
     return nargs
 end
+is_call_argument(@nospecialize(x), sv::InferenceState) =
+    isa(x, SlotNumber) && slot_id(x) ≤ narguments(sv)

base/compiler/utilities.jl

@@ -50,13 +50,23 @@ _topmod(m::Module) = ccall(:jl_base_relative_to, Any, (Any,), m)::Module
 
 function istopfunction(@nospecialize(f), name::Symbol)
     tn = typeof(f).name
-    if tn.mt.name === name
+    mn = tn.mt.name
+    if mn === name
         top = _topmod(tn.module)
         return isdefined(top, name) && isconst(top, name) && f === getglobal(top, name)
     end
     return false
 end
 
+function istoptype(@nospecialize(T), name::Symbol)
+    t = unwrap_unionall(T)
+    if isa(t, DataType) && t.name.name === name
+        top = _topmod(t.name.module)
+        return isdefined(top, name) && isconst(top, name) && T === getglobal(top, name)
+    end
+    return false
+end
+
 #######
 # AST #
 #######