irinterp: fix semi-concrete evaluation of overlay methods (#50739)

This commit introduces several fixes related to the overlay method and
its semi-concrete evaluation. When `f` is being overlayed, semi-concrete
interpretation is performed on `f`. While this isn't problematic in
itself, the issue arises since there is no overlay check within concrete
evaluation in the semi-concrete interpretation (`concrete_eval_invoke`),
potentially leading to mis-compilation. This commit makes adjustments to
allow semi-concrete interpretation for overlay methods, while preventing
the concretization of overlay methods within `concrete_eval_invoke`.

Note: Confirmed GPUCompiler test suite passes with this patch after
JuliaGPU/GPUCompiler.jl#488.

Author: aviatesk

base/compiler/abstractinterpretation.jl

@@ -28,7 +28,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
         return CallMeta(Any, Effects(), NoCallInfo())
     end
 
-    (; valid_worlds, applicable, info) = matches
+    (; valid_worlds, applicable, info, nonoverlayed) = matches
     update_valid_age!(sv, valid_worlds)
     napplicable = length(applicable)
     rettype = Bottom
@@ -39,13 +39,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
     const_results = Union{Nothing,ConstResult}[]
     multiple_matches = napplicable > 1
     fargs = arginfo.fargs
-    all_effects = EFFECTS_TOTAL
-    if !matches.nonoverlayed
-        # currently we don't have a good way to execute the overlayed method definition,
-        # so we should give up concrete eval when any of the matched methods is overlayed
-        f = nothing
-        all_effects = Effects(all_effects; nonoverlayed=false)
-    end
+    all_effects = Effects(EFFECTS_TOTAL; nonoverlayed)
 
     𝕃ₚ = ipo_lattice(interp)
     for i in 1:napplicable
@@ -792,7 +786,7 @@ function abstract_call_method_with_const_args(interp::AbstractInterpreter,
     end
     eligibility = concrete_eval_eligible(interp, f, result, arginfo, sv)
     if eligibility === :concrete_eval
-        return concrete_eval_call(interp, f, result, arginfo, sv, invokecall)
+        return concrete_eval_call(interp, f, result, arginfo, sv; invokecall)
     end
     mi = maybe_get_const_prop_profitable(interp, result, f, arginfo, si, match, sv)
     mi === nothing && return nothing
@@ -848,16 +842,17 @@ function concrete_eval_eligible(interp::AbstractInterpreter,
         add_remark!(interp, sv, "[constprop] Concrete evel disabled for inbounds")
         return :none
     end
-    if isoverlayed(method_table(interp)) && !is_nonoverlayed(effects)
-        # disable concrete-evaluation if this function call is tainted by some overlayed
-        # method since currently there is no direct way to execute overlayed methods
-        add_remark!(interp, sv, "[constprop] Concrete evel disabled for overlayed methods")
-        return :none
-    end
-    if result.edge !== nothing && is_foldable(effects)
+    mi = result.edge
+    if mi !== nothing && is_foldable(effects)
         if f !== nothing && is_all_const_arg(arginfo, #=start=#2)
-            return :concrete_eval
-        elseif !any_conditional(arginfo)
+            if is_nonoverlayed(mi.def::Method) && (!isoverlayed(method_table(interp)) || is_nonoverlayed(effects))
+                return :concrete_eval
+            end
+            # disable concrete-evaluation if this function call is tainted by some overlayed
+            # method since currently there is no easy way to execute overlayed methods
+            add_remark!(interp, sv, "[constprop] Concrete evel disabled for overlayed methods")
+        end
+        if !any_conditional(arginfo)
             return :semi_concrete_eval
         end
     end
@@ -886,8 +881,8 @@ function collect_const_args(argtypes::Vector{Any}, start::Int)
 end
 
 function concrete_eval_call(interp::AbstractInterpreter,
-    @nospecialize(f), result::MethodCallResult, arginfo::ArgInfo,
-    sv::AbsIntState, invokecall::Union{InvokeCall,Nothing})
+    @nospecialize(f), result::MethodCallResult, arginfo::ArgInfo, sv::AbsIntState;
+    invokecall::Union{InvokeCall,Nothing}=nothing)
     args = collect_const_args(arginfo, #=start=#2)
     if invokecall !== nothing
         # this call should be `invoke`d, rewrite `args` back now
@@ -1923,7 +1918,7 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn
     #     argtypes′[i] = t ⊑ a ? t : a
     # end
     𝕃ₚ = ipo_lattice(interp)
-    f = overlayed ? nothing : singleton_type(ft′)
+    f = singleton_type(ft′)
     invokecall = InvokeCall(types, lookupsig)
     const_call_result = abstract_call_method_with_const_args(interp,
         result, f, arginfo, si, match, sv, invokecall)
@@ -1934,7 +1929,7 @@ function abstract_invoke(interp::AbstractInterpreter, (; fargs, argtypes)::ArgIn
         end
     end
     rt = from_interprocedural!(interp, rt, sv, arginfo, sig)
-    effects = Effects(effects; nonoverlayed=!overlayed)
+    effects = Effects(effects; nonoverlayed = !overlayed)
     info = InvokeCallInfo(match, const_result)
     edge !== nothing && add_invoke_backedge!(sv, lookupsig, edge)
     return CallMeta(rt, effects, info)

base/compiler/methodtable.jl

@@ -167,7 +167,9 @@ end
 # This query is not cached
 findsup(@nospecialize(sig::Type), table::CachedMethodTable) = findsup(sig, table.table)
 
-isoverlayed(::MethodTableView)     = error("unsatisfied MethodTableView interface")
+isoverlayed(::MethodTableView) = error("unsatisfied MethodTableView interface")
 isoverlayed(::InternalMethodTable) = false
-isoverlayed(::OverlayMethodTable)  = true
+isoverlayed(::OverlayMethodTable) = true
 isoverlayed(mt::CachedMethodTable) = isoverlayed(mt.table)
+isoverlayed(m::Method) = isdefined(m, :external_mt)
+is_nonoverlayed(m::Method) = !isoverlayed(m)

base/compiler/ssair/irinterp.jl

@@ -14,7 +14,8 @@ function concrete_eval_invoke(interp::AbstractInterpreter,
     argtypes = collect_argtypes(interp, inst.args[2:end], nothing, irsv)
     argtypes === nothing && return Pair{Any,Bool}(Bottom, false)
     effects = decode_effects(code.ipo_purity_bits)
-    if is_foldable(effects) && is_all_const_arg(argtypes, #=start=#1)
+    if (is_foldable(effects) && is_all_const_arg(argtypes, #=start=#1) &&
+        is_nonoverlayed(effects) && is_nonoverlayed(mi.def::Method))
         args = collect_const_args(argtypes, #=start=#1)
         value = let world = get_world_counter(interp)
             try

test/compiler/AbstractInterpreter.jl

@@ -9,10 +9,10 @@ include("newinterp.jl")
 # OverlayMethodTable
 # ==================
 
-import Base.Experimental: @MethodTable, @overlay
+using Base.Experimental: @MethodTable, @overlay
 
 @newinterp MTOverlayInterp
-@MethodTable(OverlayedMT)
+@MethodTable OverlayedMT
 CC.method_table(interp::MTOverlayInterp) = CC.OverlayMethodTable(CC.get_world_counter(interp), OverlayedMT)
 
 function CC.add_remark!(interp::MTOverlayInterp, ::CC.InferenceState, remark)
@@ -113,12 +113,50 @@ end |> only === Nothing
 @MethodTable Issue48097MT
 CC.method_table(interp::Issue48097Interp) = CC.OverlayMethodTable(CC.get_world_counter(interp), Issue48097MT)
 CC.InferenceParams(::Issue48097Interp) = CC.InferenceParams(; unoptimize_throw_blocks=false)
+function CC.concrete_eval_eligible(interp::Issue48097Interp,
+    @nospecialize(f), result::CC.MethodCallResult, arginfo::CC.ArgInfo, sv::CC.AbsIntState)
+    ret = @invoke CC.concrete_eval_eligible(interp::CC.AbstractInterpreter,
+        f::Any, result::CC.MethodCallResult, arginfo::CC.ArgInfo, sv::CC.AbsIntState)
+    if ret === :semi_concrete_eval
+        # disable semi-concrete interpretation
+        return :none
+    end
+    return ret
+end
 @overlay Issue48097MT @noinline Core.throw_inexacterror(f::Symbol, ::Type{T}, val) where {T} = return
 issue48097(; kwargs...) = return 42
 @test fully_eliminated(; interp=Issue48097Interp(), retval=42) do
     issue48097(; a=1f0, b=1.0)
 end
 
+# Should not concrete-eval overlayed methods in semi-concrete interpretation
+@newinterp OverlaySinInterp
+@MethodTable OverlaySinMT
+CC.method_table(interp::OverlaySinInterp) = CC.OverlayMethodTable(CC.get_world_counter(interp), OverlaySinMT)
+overlay_sin1(x) = error("Not supposed to be called.")
+@overlay OverlaySinMT overlay_sin1(x) = cos(x)
+@overlay OverlaySinMT Base.sin(x::Union{Float32,Float64}) = overlay_sin1(x)
+let oc = Base.code_ircode(; interp=OverlaySinInterp()) do
+        sin(0.)
+    end |> only |> first |> Core.OpaqueClosure
+    @test oc() == cos(0.)
+end
+@overlay OverlaySinMT Base.sin(x::Union{Float32,Float64}) = @noinline overlay_sin1(x)
+let oc = Base.code_ircode(; interp=OverlaySinInterp()) do
+        sin(0.)
+    end |> only |> first |> Core.OpaqueClosure
+    @test oc() == cos(0.)
+end
+_overlay_sin2(x) = error("Not supposed to be called.")
+@overlay OverlaySinMT _overlay_sin2(x) = cos(x)
+overlay_sin2(x) = _overlay_sin2(x)
+@overlay OverlaySinMT Base.sin(x::Union{Float32,Float64}) = @noinline overlay_sin2(x)
+let oc = Base.code_ircode(; interp=OverlaySinInterp()) do
+        sin(0.)
+    end |> only |> first |> Core.OpaqueClosure
+    @test oc() == cos(0.)
+end
+
 # AbstractLattice
 # ===============