From 38d051dec2b09bdca00f6fae3fc806e677bd14f6 Mon Sep 17 00:00:00 2001 From: Shuhei Kadowaki Date: Wed, 14 Sep 2022 01:41:16 +0900 Subject: [PATCH] inference: concretize `invoke` callsite correctly It turns out that previously we didn't concretize `invoke` callsite correctly, as we didn't take into account whether or not the call being concretized is `invoke`-d or not, e.g.: ``` julia> invoke_concretized2(a::Int) = a > 0 ? :int : nothing invoke_concretized2 (generic function with 1 method) julia> invoke_concretized2(a::Integer) = a > 0 ? :integer : nothing invoke_concretized2 (generic function with 2 methods) julia> let Base.Experimental.@force_compile Base.@invoke invoke_concretized2(42::Integer) end :int # this should return `:integer` instead ``` This commit fixes that up by propagating information `invoke`-d callsite to `concrete_eval_call`. Now we should pass the following test cases: ```julia invoke_concretized1(a::Int) = a > 0 ? :int : nothing invoke_concretized1(a::Integer) = a > 0 ? "integer" : nothing @test Base.infer_effects((Int,)) do a @invoke invoke_concretized1(a::Integer) end |> Core.Compiler.is_foldable @test Base.return_types() do @invoke invoke_concretized1(42::Integer) end |> only === String invoke_concretized2(a::Int) = a > 0 ? :int : nothing invoke_concretized2(a::Integer) = a > 0 ? :integer : nothing @test Base.infer_effects((Int,)) do a @invoke invoke_concretized2(a::Integer) end |> Core.Compiler.is_foldable @test let Base.Experimental.@force_compile @invoke invoke_concretized2(42::Integer) end === :integer ``` --- base/compiler/abstractinterpretation.jl | 24 ++++++++++++++++++++---- test/compiler/inference.jl | 21 +++++++++++++++++++++ 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl index 7f2106b8f0070..9eb764848eb12 100644 --- a/base/compiler/abstractinterpretation.jl +++ b/base/compiler/abstractinterpretation.jl @@ -885,6 +885,7 @@ function abstract_call_method_with_const_args(interp::AbstractInterpreter, resul add_backedge!(res.const_result.mi, sv, invoketypes) return res end + f isa InvokeCall && (f = f.f) # unwrap `invoke`-ed function (the call will be virtualized) mi = maybe_get_const_prop_profitable(interp, result, f, arginfo, match, sv) mi === nothing && return nothing # try semi-concrete evaluation @@ -1674,10 +1675,10 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn nargtype isa DataType || return CallMeta(Any, Effects(), false) # other cases are not implemented below isdispatchelem(ft) || return CallMeta(Any, Effects(), false) # check that we might not have a subtype of `ft` at runtime, before doing supertype lookup below ft = ft::DataType - types = rewrap_unionall(Tuple{ft, unwrap_unionall(types).parameters...}, types)::Type + lookupsig = rewrap_unionall(Tuple{ft, unwrap_unionall(types).parameters...}, types)::Type nargtype = Tuple{ft, nargtype.parameters...} argtype = Tuple{ft, argtype.parameters...} - match, valid_worlds, overlayed = findsup(types, method_table(interp)) + match, valid_worlds, overlayed = findsup(lookupsig, method_table(interp)) match === nothing && return CallMeta(Any, Effects(), false) update_valid_age!(sv, valid_worlds) method = match.method @@ -1685,7 +1686,7 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn ti = tienv[1]; env = tienv[2]::SimpleVector result = abstract_call_method(interp, method, ti, env, false, sv) (; rt, edge, effects) = result - edge !== nothing && add_backedge!(edge::MethodInstance, sv, types) + edge !== nothing && add_backedge!(edge::MethodInstance, sv, lookupsig) match = MethodMatch(ti, env, method, argtype <: method.sig) res = nothing sig = match.spec_types @@ -1697,8 +1698,16 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn # t, a = ti.parameters[i], argtypes′[i] # argtypes′[i] = t ⊑ a ? t : a # end + # form `InvokeCall` for possible concretization + invokef = nothing + if !overlayed + f = singleton_type(ft′) + if f !== nothing + invokef = InvokeCall(f, types) + end + end const_call_result = abstract_call_method_with_const_args(interp, result, - overlayed ? nothing : singleton_type(ft′), arginfo, match, sv, types) + invokef, arginfo, match, sv, lookupsig) const_result = nothing if const_call_result !== nothing if ⊑(typeinf_lattice(interp), const_call_result.rt, rt) @@ -1709,6 +1718,13 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn return CallMeta(from_interprocedural!(ipo_lattice(interp), rt, sv, arginfo, sig), effects, InvokeCallInfo(match, const_result)) end +struct InvokeCall # used for concrete evaluation + f + types # ::Type + InvokeCall(@nospecialize(f), @nospecialize(types)) = new(f, types) +end +(invokef::InvokeCall)(@nospecialize args...) = invoke(invokef.f, invokef.types, args...) + function invoke_rewrite(xs::Vector{Any}) x0 = xs[2] newxs = xs[3:end] diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl index e0b07422f3821..ed7dcbdfb0fd0 100644 --- a/test/compiler/inference.jl +++ b/test/compiler/inference.jl @@ -4211,3 +4211,24 @@ function isa_kindtype(T::Type{<:AbstractVector}) return nothing end @test only(Base.return_types(isa_kindtype)) === Union{Nothing,Symbol} + +invoke_concretized1(a::Int) = a > 0 ? :int : nothing +invoke_concretized1(a::Integer) = a > 0 ? "integer" : nothing +# check if `invoke(invoke_concretized1, Tuple{Integer}, ::Int)` is foldable +@test Base.infer_effects((Int,)) do a + @invoke invoke_concretized1(a::Integer) +end |> Core.Compiler.is_foldable +@test Base.return_types() do + @invoke invoke_concretized1(42::Integer) +end |> only === String + +invoke_concretized2(a::Int) = a > 0 ? :int : nothing +invoke_concretized2(a::Integer) = a > 0 ? :integer : nothing +# check if `invoke(invoke_concretized2, Tuple{Integer}, ::Int)` is foldable +@test Base.infer_effects((Int,)) do a + @invoke invoke_concretized2(a::Integer) +end |> Core.Compiler.is_foldable +@test let + Base.Experimental.@force_compile + @invoke invoke_concretized2(42::Integer) +end === :integer