Give the default method table view one type for every stack depth

src/interface.jl

@@ -680,6 +680,8 @@ from [`method_tables`](@ref), then [`SHARED_METHOD_TABLE`](@ref), over Base's me
 
 Back-ends should not need to override this. Ones that do replace the whole stack, including
 `SHARED_METHOD_TABLE`, so code compiled for them behaves as if it were compiled by Julia.
+Returning a view of another type also gives inference a different interpreter type, so the
+back-end doesn't share the inference code that GPUCompiler precompiles.
 """
 method_table_view(@nospecialize(job::CompilerJob)) =
     stack_method_tables(job.world, method_tables(job)..., SHARED_METHOD_TABLE)

src/jlgen.jl

@@ -156,11 +156,99 @@ Base.Experimental.@MethodTable(SHARED_METHOD_TABLE)
     using CompilerCaching: StackedMethodTable
 end
 
-# look up methods in each table in turn, then in Base
-stack_method_tables(world::UInt) = CC.InternalMethodTable(world)
-stack_method_tables(world::UInt, mt::Core.MethodTable) = StackedMethodTable(world, mt)
-stack_method_tables(world::UInt, mt::Core.MethodTable, rest::Core.MethodTable...) =
-    StackedMethodTable(world, mt, stack_method_tables(world, rest...))
+# Look up methods in each table in turn, then in Base. Unlike nesting `StackedMethodTable`s,
+# this has the same type for every stack depth, so `GPUInterpreter` has one concrete type for
+# all back-ends using the default `method_table_view`, and the inference specializations
+# precompiled in GPUCompiler's own image are reused by every back-end.
+struct MethodTableStack <: CC.MethodTableView
+    world::UInt
+    tables::Vector{Core.MethodTable}
+end
+
+CC.isoverlayed(table::MethodTableStack) = !isempty(table.tables)
+
+@static if VERSION >= v"1.14.0-DEV.3088"
+    stack_findall(@nospecialize(sig::Type), mt::Union{Nothing,Core.MethodTable}, world::UInt,
+                  limit::Int, include_ambiguous::Bool) =
+        CC._findall(sig, mt, world, limit, include_ambiguous)
+else
+    stack_findall(@nospecialize(sig::Type), mt::Union{Nothing,Core.MethodTable}, world::UInt,
+                  limit::Int, ::Nothing) =
+        CC._findall(sig, mt, world, limit)
+end
+
+# Equivalent to the nested `StackedMethodTable` lookup: accumulate matches from each table
+# until one fully covers `sig`, falling back to Base if none does.
+function stack_findall(@nospecialize(sig::Type), table::MethodTableStack, limit::Int,
+                       include_ambiguous)
+    acc = nothing
+    for mt in table.tables
+        result = stack_findall(sig, mt, table.world, limit, include_ambiguous)
+        result === nothing && return nothing # too many matches
+        acc = acc === nothing ? result : merge_lookup(acc, result)
+        nr = CC.length(result)
+        if nr >= 1 && CC.getindex(result, nr).fully_covers
+            return acc, true
+        end
+    end
+    result = stack_findall(sig, nothing, table.world, limit, include_ambiguous)
+    result === nothing && return nothing
+    overlayed = acc !== nothing && !CC.isempty(acc)
+    return (acc === nothing ? result : merge_lookup(acc, result)), overlayed
+end
+
+merge_lookup(a::CC.MethodLookupResult, b::CC.MethodLookupResult) =
+    CC.MethodLookupResult(
+        CC.vcat(a.matches, b.matches),
+        CC.WorldRange(CC.max(a.valid_worlds.min_world, b.valid_worlds.min_world),
+                      CC.min(a.valid_worlds.max_world, b.valid_worlds.max_world)),
+        a.ambig | b.ambig)
+
+@static if VERSION >= v"1.14.0-DEV.3088"
+    function CC.findall(@nospecialize(sig::Type), table::MethodTableStack;
+                        limit::Int=-1, include_ambiguous::Bool=false)
+        r = stack_findall(sig, table, limit, include_ambiguous)
+        r === nothing ? nothing : r[1]
+    end
+elseif HAS_INTEGRATED_CACHE
+    function CC.findall(@nospecialize(sig::Type), table::MethodTableStack; limit::Int=-1)
+        r = stack_findall(sig, table, limit, nothing)
+        r === nothing ? nothing : r[1]
+    end
+else
+    # 1.10 wraps the lookup in a `MethodMatchResult` that records overlay use
+    function CC.findall(@nospecialize(sig::Type), table::MethodTableStack; limit::Int=-1)
+        r = stack_findall(sig, table, limit, nothing)
+        r === nothing ? nothing : CC.MethodMatchResult(r[1], r[2])
+    end
+end
+
+function CC.findsup(@nospecialize(sig::Type), table::MethodTableStack)
+    min_world, max_world = typemin(UInt), typemax(UInt)
+    for mt in table.tables
+        match, valid_worlds = CC._findsup(sig, mt, table.world)
+        min_world = max(min_world, valid_worlds.min_world)
+        max_world = min(max_world, valid_worlds.max_world)
+        if match !== nothing
+            @static if HAS_INTEGRATED_CACHE
+                return match, CC.WorldRange(min_world, max_world)
+            else
+                return match, CC.WorldRange(min_world, max_world), true
+            end
+        end
+    end
+    match, valid_worlds = CC._findsup(sig, nothing, table.world)
+    min_world = max(min_world, valid_worlds.min_world)
+    max_world = min(max_world, valid_worlds.max_world)
+    @static if HAS_INTEGRATED_CACHE
+        return match, CC.WorldRange(min_world, max_world)
+    else
+        return match, CC.WorldRange(min_world, max_world), false
+    end
+end
+
+stack_method_tables(world::UInt, mts::Core.MethodTable...) =
+    MethodTableStack(world, Core.MethodTable[mts...])
 
 
 ## interpreter

test/native.jl

@@ -2862,6 +2862,19 @@ end
     end
 end
 
+@testset "interpreter type doesn't depend on the method tables" begin
+    # so that the inference code that GPUCompiler precompiles serves every back-end
+    mod = @eval module $(gensym())
+        Base.Experimental.@MethodTable(method_table)
+        Base.Experimental.@MethodTable(other_method_table)
+    end
+    interpreter_type(method_table) =
+        typeof(GPUCompiler.get_interpreter(Native.create_job(identity, (Int,); method_table)[1]))
+    @test interpreter_type(mod.method_table) ==
+          interpreter_type((mod.method_table, mod.other_method_table)) ==
+          interpreter_type((mod.method_table, mod.other_method_table, mod.method_table))
+end
+
 @testset "runtime functions from overlay methods" begin
     # runtime library functions (`signal_exception`, `malloc`, ...) should be resolved
     # through the job's method table, so that back-ends can keep GPU-only code out of

test/utils.jl

@@ -352,6 +352,58 @@ next_world = Base.get_world_counter()
 end
 
 
+module StackTables
+    Base.Experimental.@MethodTable(A)
+    Base.Experimental.@MethodTable(B)
+    Base.Experimental.@MethodTable(C)
+    f(x) = 0
+    Base.Experimental.@overlay A f(x::Int) = 1
+    Base.Experimental.@overlay B f(x::Integer) = 2
+    Base.Experimental.@overlay C f(x) = 3
+    g(x) = 0
+    Base.Experimental.@overlay B g(x::Float64) = 1
+    h(x::Int) = 0
+    h(x::Float64) = 0
+    function k end
+    Base.Experimental.@overlay C k(x) = 1
+    function amb end
+    Base.Experimental.@overlay A amb(x::Int, y) = 1
+    Base.Experimental.@overlay A amb(x, y::Int) = 2
+end
+
+@testset "stacked method tables" begin
+    # the flat stack that `method_table_view` uses behaves like nested `StackedMethodTable`s
+    world = Base.get_world_counter()
+    T = StackTables
+    lookup(r) = r isa Core.Compiler.MethodLookupResult ?
+        ([(m.method, m.spec_types, m.fully_covers) for m in r.matches],
+         r.valid_worlds.min_world, r.valid_worlds.max_world, r.ambig) :
+        r === nothing ? nothing : (lookup(r.matches), r.overlayed)
+    sup((match, worlds, rest...)) =
+        (match === nothing ? nothing : match.method, worlds.min_world, worlds.max_world, rest...)
+    for (tables, nested) in (
+            ((T.A, T.B, T.C), StackedMethodTable(world, T.A,
+                                  StackedMethodTable(world, T.B, StackedMethodTable(world, T.C)))),
+            ((T.A, T.B), StackedMethodTable(world, T.A, StackedMethodTable(world, T.B))))
+        flat = GPUCompiler.stack_method_tables(world, tables...)
+        @test isoverlayed(flat)
+        for sig in (Tuple{typeof(T.f),Int}, Tuple{typeof(T.f),Int32}, Tuple{typeof(T.f),Any},
+                    Tuple{typeof(T.f),Float64}, Tuple{typeof(T.g),Any},
+                    Tuple{typeof(T.g),Float64}, Tuple{typeof(T.h),Any}, Tuple{typeof(T.h),Int},
+                    Tuple{typeof(T.k),Int}, Tuple{typeof(T.amb),Int,Int},
+                    Tuple{typeof(T.amb),Any,Any}, Tuple{typeof(+),Any,Any}),
+            limit in (-1, 1, 3)
+            @test lookup(findall(sig, flat; limit)) == lookup(findall(sig, nested; limit))
+            @test sup(findsup(sig, flat)) == sup(findsup(sig, nested))
+        end
+    end
+    @test !isoverlayed(GPUCompiler.stack_method_tables(world))
+
+    # the comparison above covers `ambig`
+    r = findall(Tuple{typeof(T.amb),Any,Any}, GPUCompiler.stack_method_tables(world, T.A))
+    @test (r isa Core.Compiler.MethodLookupResult ? r : r.matches).ambig
+end
+
 @testset "Mock Enzyme" begin
     id = Enzyme.deferred_codegen_id(typeof(identity), Tuple{Vector{Float64}}, Val(false))
     @test id isa UInt && id >= UInt(1) << 62