Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add support for external method tables (JuliaLang#39697)
This PR implements a way to keep tables of methods that are not part of the internal method table, but still participate in the special support we have for keeping tables of methods, in particular unification through precompilation and efficient lookup. The intended design use case is to allow for method overlay tables for various non-CPU backends (e.g. GPU and TPU). These backends would like to modify basic function like `sin` to perform better on the device in question (or in the case of TPU to define them over non-LLVM intrinsics). To date, the best available mechanism of achieving this result was to use a Cassette-like pass rewriting every method and injecting an overlay if necessary. However, this approach is somewhat unsatisfying for two reasons: 1. It requires rewriting every function, which has non-trivial performance cost. 2. It is (currently) expensive because of the repeated calls to generated functions. 3. It confuses inference, because suddenly everything is one method. We have hooks to work around this, but support is incomplete. It is also not clear that Cassette it is the best conceptual model, because these methods *are* methods of the same generic function, they just happen to only be applicable for a particular backend. It is worth noting that this PR only gives the ability to keep these tables of methods. It assigns no particular meaning to them and the runtime (and regular inference) do not look at them. They are designed as an implementation detail for external compilers and similar tools. This feature does not replace Cassette for the method-interception use case in the absence of such a compiler, though it could in the future form part of a solution (I'm hoping the AD work will in due course lead to abstractions that would enable a "faster Cassette" which may use part of these fetaures). As such, I'm not sure we'll ever move this out of Experimental, but until such a time that we have a better solution, I think this'll be a useful feature for the GPU stack. With all those disclaimers out of the way, here is a description of the various parts of the current design that deserve discussion: # Demo ```julia julia> using Base.Experimental: @overlay, @MethodTable julia> @MethodTable(mt) # 0 methods: julia> mt # 0 methods: julia> @overlay mt function sin(x::Float64) 1 end julia> @overlay mt function cos(x::Float64) 1 end julia> mt # 2 methods: [1] cos(x::Float64) in Main at REPL[5]:1 [2] sin(x::Float64) in Main at REPL[4]:1 julia> Base._methods_by_ftype(Tuple{typeof(sin), Float64}, mt, 1, typemax(UInt)) 1-element Vector{Any}: Core.MethodMatch(Tuple{typeof(sin), Float64}, svec(), sin(x::Float64) in Main at REPL[4]:1, true) julia> Base._methods_by_ftype(Tuple{typeof(sin), Float64}, 1, typemax(UInt)) 1-element Vector{Any}: Core.MethodMatch(Tuple{typeof(sin), Float64}, svec(Float64), sin(x::T) where T<:Union{Float32, Float64} in Base.Math at special/trig.jl:29, true) ``` # The `@overlay` macro The macro replaces the function name by an `Expr(:overlay, mt, name)`, which then gets piped through to Method def. One particular design aspect here is that I've stopped letting the 4-argument :method Expr introduce new generic functions, reserving this functionality entirely to the 2-argument :method Expr. We already started going this way when we began omitting method names from the 4-argument version. This PR re-uses that name field of the 4-argument version to specify a method table instead. # Identity of methods I think one of the biggest questions of this design is what happens to the identity of methods. Until OpaqueClosure, all methods were uniquely identified by their signatures, with the applicable method table computed from the first argument of the signature. This is important so that incremental compilation can properly merge method tables coming from different .ji files. For these methods, that is of course not the correct method table to use for these methods, so methods that are not part of the internal method table will instead have a backreference to the applicable method table. # Identity of method tables Method tables are identified by the name of their binding in the containing module. To ensure consistency of this mapping, these MethodTables may only be constructed using the `@MethodTable(name)` macro, which simultaneously establishes a const binding in the declaring module. Co-authored-by: Tim Besard <[email protected]>
- Loading branch information
