Traitor.jl icon indicating copy to clipboard operation
Traitor.jl copied to clipboard

Get this working again, add betray!

Open MasonProtter opened this issue 4 years ago • 7 comments

So I've gotten the package working on julia version 1.4+, and I've got a working implementation of the betray! functionality.

One big annoyance that's still present is worldage issues:

julia> using Traitor

julia> begin
           abstract type Size end
           struct Small  <: Size end
           Size(::Union{Type{Int32},Type{Int64}}) = Small
       end
Size

julia> @traitor f(x::::Small) = x + 1
_f{::Any::Tuple{Small}} (generic function with 1 method)

julia> f(1)
2

julia> Size(::Type{Float64}) = Small
Size

julia> f(1.0)
ERROR: MethodError: no method matching Size(::Type{Float64})
The applicable method may be too new: running in world age 29605, while current world is 29607.
Closest candidates are:
  Size(::Type{Float64}) at REPL[5]:1 (method too new to be called from this world context.)
  Size(::Union{Type{Int32}, Type{Int64}}) at REPL[2]:4
Stacktrace:
[...]

However, I believe this can be solved through the injection of backedges into the generated function body. I think the above worldage issue would be solved by just copying the backedges from Traitor.trait_dispatch into the code info for the generated function returned by @traitor.

Another approach that might help us get away from generated functions would be to use https://github.com/RelationalAI-oss/Salsa.jl rather than generated functions to cache, memoize, and reactively update the dispatch lookup tables.

Yet another option would be to use the AbstractInterpreter machinery. I'm really excited about this machinery and would like to learn to use it, but it's still pretty intimidating to me and looks quite difficult to learn to use so I haven't gotten around to it yet.

MasonProtter avatar Feb 18 '21 19:02 MasonProtter

CC @timholy if case you have any thoughts or comments.

MasonProtter avatar Feb 18 '21 20:02 MasonProtter

However, I believe this can be solved through the injection of backedges into the generated function body. I think the above worldage issue would be solved by just copying the backedges from Traitor.trait_dispatch into the code info for the generated function returned by @traitor.

Yes, I think this is the approach we should take. The idea here was to have fully-compiled code (rather than memoization like Salsa.jl) and also to make a proof-of-principle that could potentially (if the core devs were interested) be integrated into Julia proper.

Hm, from that POV then I think the AbstractInterpreter machinery would be the way to go because it'd be using the same API that julia already uses for compilation and deciding dispatch. I can't imagine tje Compiler team being at all enthusiastic about a PR to Core.Compiler that uses generated functions to decide dispatch.

Either way though, I'll get the backedge injection route working first because I understand it better.

MasonProtter avatar Feb 19 '21 01:02 MasonProtter

There are a couple of packages that subtype AbstractInterpreter to good effect and your interest in the same had me ask Keno how they'd play together. Turns out they don't. Subtyping AI is apparently a temporary hack until the higher level compiler interfaces can get figured out. Obviously the ideal route is through compiler passes, but given that information (this may clash or have undefined behavior with GPUCompiler etc) not sure if that changes how you want to proceed.

@0x0f0f0f 's metatheory.jl is one candidate for this higher level compiler infra, which might shift to integrate with SSA IR.

AriMKatz avatar Feb 19 '21 01:02 AriMKatz

relevant https://github.com/JuliaLang/julia/pull/39697

AriMKatz avatar Feb 19 '21 01:02 AriMKatz

Okay, so the good news is that I managed to get the invalidation machinery to cooperate and recompile stuff as needed:

julia> using Traitor

julia> begin
           abstract type Size end
           struct Small  <: Size end
           Size(::Union{Type{Int32},Type{Int64}}) = Small
       end
Size

julia> @traitor f(x::::Small) = x + 1
f (generic function with 1 method)

julia> f(1)
2

julia> Size(::Type{Float64}) = Small
Size

julia> f(1.0)
2.0

The bad new is that I can't get it to work without a Cassette overdub pass through trait_dispatch that doesn't actually do anything. I was debugging stuff with a cassette pass using some code from https://github.com/NHDaly/StagedFunctions.jl/, and suddenly I got it to work. Without the pass, it wouldn't work. Then I made the pass literally do nothing and it still worked.

I suspect that maybe the pass is forcing some specializations that the compiler was giving up on or something. I'm not sure and I've given up for now on figuring it out. So for now, my branch depends on Cassette.jl.

This is weird.

@NHDaly, you wouldn't happen to know what's going on here, would you? The code in https://github.com/andyferris/Traitor.jl/blob/98ded521ca8492860818457b9de81343b3fd9fa2/src/Traitor.jl#L279-L356 is essentially a pruned version of what's in StagedFunctions.jl

MasonProtter avatar Feb 19 '21 04:02 MasonProtter

Kinda swamped but just chiming in to say exciting stuff! @andyferris is more than capable of handling this but I will look sometime (before or after merging).

timholy avatar Feb 19 '21 21:02 timholy

Not a problem Tim, I know you're a busy guy!

MasonProtter avatar Feb 19 '21 21:02 MasonProtter