External Method Tables

[Keno]

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> 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.

TODO:

• Core.Compiler interfaces to make this easy to use with
  external AbstractInterpreters
• Serialization.jl support
• Tests
• Docs

[timholy]

More excitement for Revise! 😆

But this makes a lot of sense, so 👍

[maleadt]

@overlay can currently be used with functions in other modules without having to explicitly import/specify them (i.e. @overlay sin instead of @overlay Base.sin). That makes sense, as it can't be used to introduce new functions -- it should always refer to an existing one -- but maybe it should behave similar to extending functions? If so, where's the logic for that?

Another design consideration: the current design doesn't allow adding methods to another module's overlay method table, or at least that won't survive precompilation. Maybe that's acceptable though?

[Keno]

Another design consideration: the current design doesn't allow adding methods to another module's overlay method table, or at least that won't survive precompilation. Maybe that's acceptable though?

This was supposed to work. I added extra code in serialization for it. I may have done it wrong though.

[maleadt]

This was supposed to work. I added extra code in serialization for it. I may have done it wrong though.

OK, good to know, I'll debug it then.

[maleadt]

I added extra code in serialization for it.

Where's that code? I only see the changes to refer to another module's method table when (de)serializing a method, but nothing to save or merge another module's method table? And just to be clear:

Foo.jl:

module Foo
  Base.Experimental.@MethodTable(mt)
end

Bar.jl:

module Bar
  using Foo
  Base.Experimental.@overlay Foo.mt sin(x::Int) = 1
end

using Foo, Bar
@show Foo.mt

Foo.mt = # 0 methods

[maleadt]

@Keno Thoughts on #39697 (comment)?
Also, what needs to happen for Serialization.jl? serialize_any seems to work fine for method tables; I assume you want it to maintain the same contract @MethodTable enforces (serialize the owning module and binding name, and assign it like that)? That seems like weird behaviour for Serialization's API though...

[maleadt]

Ah, I forgot we hadn't merged this yet. @Keno, do you want to fold this in your upcoming compiler plugin work, or should we merge this first?

[Keno]

I'm happy merging this as is and revising it later if necessary.