From 6eb6bacc4008ae63f1f7791775c5de3d71f658b9 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 5 Nov 2026 20:04:46 +0100 Subject: [PATCH] Give the default method table view one type for every stack depth The default `method_table_view` nested a `StackedMethodTable` per method table, and `GPUInterpreter` is parametric in the view, so back-ends stacking a different number of tables got different interpreter types. The inference code specialized on that type is only shared between packages whose stacks happen to have the same depth: GPUCompiler's own precompilation covers none of the back-ends, and Metal.jl precompiled ~15 s of CPU time faster only because KernelAbstractions' POCL back-end happened to stack as many tables, until KernelAbstractions 0.10.5 added one. Look up the tables in a flat `MethodTableStack` instead, with the same semantics as the nested chain, so every back-end that returns its tables from `method_tables` shares one interpreter type, and the inference code that GPUCompiler precompiles. --- src/interface.jl | 2 + src/jlgen.jl | 98 +++++++++++++++++++++++++++++++++++++++++++++--- test/native.jl | 13 +++++++ test/utils.jl | 52 +++++++++++++++++++++++++ 4 files changed, 160 insertions(+), 5 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index 00dcfea3..419d66b2 100644 --- a/src/interface.jl +++ b/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) diff --git a/src/jlgen.jl b/src/jlgen.jl index 940c64e3..33f77f13 100644 --- a/src/jlgen.jl +++ b/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 diff --git a/test/native.jl b/test/native.jl index 8bf1008a..e7959e0c 100644 --- a/test/native.jl +++ b/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 diff --git a/test/utils.jl b/test/utils.jl index 95b07a81..6e418111 100644 --- a/test/utils.jl +++ b/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