generators: expose caller's world to GeneratedFunctionStub

[vtjnash]

This has been a very long request, particularly from GPU folks, so tagging them here to be aware of this change. In particular, see the changes to some of the representative tests (e.g. compiler/context.jl) to understand how to update your code to work with this and now support #265.

Expose the demanded world to the GeneratedFunctionStub caller, for users such as Cassette. If this argument is used, the uesr must return a CodeInfo with the min/max world field set correctly.

Make the internal representation a tiny bit more compact also, removing a little bit of unnecessary metadata.

Remove support for returning body isa CodeInfo via this wrapper, since it is impossible to return a correct object via the GeneratedFunctionStub since it strips off the world argument, which is required for it to do so. This also removes support for not inferring these fully (expand_early=false).

Also answer method lookup queries about the future correctly, by refusing to answer them. This helps keeps execution correct as methods get added to the system asynchronously.

This reverts "fix # 25678: return matters for generated functions (# 40778)", since this is no longer sensible to return here anyways, so it is no longer permitted or supported by this macro.

Fixes various issues where we failed to specify the correct world.

[maleadt]

Thanks, I'll give this a proper look as soon as I'm done fixing GPU packages for 1.9.

[jpsamaroo]

This is not my area of expertise, so I'll defer to @maleadt 's judgement

[maleadt]

Thanks for this. I've adapted the GPUCompiler.jl code that uses this, and everything still works, so no functionality is lost.

I however wasn't able to use the feature for what I expected it to be used for, i.e., being able to use the correct world where we now always use typemax(UInt). Here's a small demo:

using Core.Compiler: retrieve_code_info, CodeInfo, MethodInstance, ReturnNode
using Base: _methods_by_ftype


# generated function that returns the world in which we should compile a function.
# it also registers backedges so that it gets invalidated on redefinitions.

function get_world_generator(world::UInt, source, self, f, tt)
    @nospecialize
    Core.println("generator running for world ", Int(world))

    # get a hold of the method and code info of the kernel function
    sig = Tuple{f, tt.parameters...}
    mthds = _methods_by_ftype(sig, -1, world)
    mtypes, msp, m = mthds[1]
    mi = ccall(:jl_specializations_get_linfo, Ref{MethodInstance}, (Any, Any, Any), m, mtypes, msp)
    ci = retrieve_code_info(mi, world)::CodeInfo

    # prepare a new code info
    new_ci = copy(ci)
    empty!(new_ci.code)
    empty!(new_ci.codelocs)
    resize!(new_ci.linetable, 1)
    empty!(new_ci.ssaflags)
    new_ci.ssavaluetypes = 0
    new_ci.edges = MethodInstance[mi]

    # prepare the slots
    new_ci.slotnames = Symbol[Symbol("#self#"), :f, :tt]
    new_ci.slotflags = UInt8[0x00 for i = 1:3]

    # return the world
    push!(new_ci.code, ReturnNode(Int(world)))
    push!(new_ci.ssaflags, 0x00)
    push!(new_ci.codelocs, 1)
    new_ci.ssavaluetypes += 1

    return new_ci
end

@eval function get_world(f, tt)
    $(Expr(:meta, :generated_only))
    $(Expr(:meta, :generated, get_world_generator))
end


## main

println("defining kernel")
kernel() = 1
println("toplevel got world ", get_world(kernel, ()))
println()

println("spawning task")
c1, c2 = Condition(), Condition()
function background(kernel, ())
    println("task running in world ", ccall(:jl_get_tls_world_age, Int, ()))
    notify(c1)
    wait(c2)    # wait for redefinition
    println("task calling generator from world ", ccall(:jl_get_tls_world_age, UInt, ()))
    world = get_world(kernel, ())
    println("task got world ", world)
end
t = @async background(kernel, ())
wait(c1)
println()

println("redefine kernel")
kernel() = 2
println("world is now ", Base.get_world_counter())
println("toplevel got world ", get_world(kernel, ()))
println()

println("waiting for task")
notify(c2)  # wake up the task
wait(t)

In summary, I define a kernel, launch a task that uses it, and then redefine the kernel. Here, 'using the kernel' is just returning the world age that the generator passes us; in reality we'd be using that to index a compilation cache. I would have expected the world that gets returned inside the task to be less than the world after the redefinition, but instead world is more recent and as a result the task gets to see the redefinition:

defining kernel
generator running for world 1
toplevel got world 1

spawning task
generator running for world 2
task running in world 2

redefine kernel
world is now 3
generator running for world 3
toplevel got world 3

waiting for task
task calling generator from world 2
task got world 3

I'm not sure why the task got world 3, as the generator that ran after the @async definition was using world 2... Am I misunderstanding how this is supposed to work, and if so, any thoughts on how we should make this work?

[vtjnash]

I don't think you are setting world bounds and edges on the result, which are mandatory fields, or it is just undefined entirely what you get (I guess this could be enforced better?)

[maleadt]

Ah, OK. Setting new_ci.min_world and max_world does indeed seem to help (I assumed they would be properly set in the ci I retrieved), but I need to do some more testing. Edges I was already setting, or is there something I should set?

[vtjnash]

Ah, I see the edges now. The _methods_by_ftype call also (optionally) returns age bounds, which you are required to propagate (usually pretty easy, since it is just min/max calls. ci does not normally have world restrictions, just method lookup is what is adding restrictions to it

[maleadt]

OK thanks. I'm still in the process of adapting, but this looks good to me.

EDIT: Rebased.

[maleadt]

Let's do a PkgEval so that we know which packages to fix:

@nanosoldier runtests()

[nanosoldier]

Your package evaluation job has completed - possible new issues were detected.
A full report can be found here.

[maleadt]

Well, that's significantly more failures than I assumed there would be (i.e., more than just those packages using GeneratedFunctionStub). Looks like ConstructionBase.jl is part of many of those failures; I'll have a look.

[maleadt]

@vtjnash Is it an intended side-effect of this PR that you can't call which anymore from a generator?

@generated function test()
    Core.println(which(identity, Tuple{Nothing}))
    return :()
end
test()

❯ julia +1.9 wip.jl
identity(Any)

❯ julia +dev wip.jl
ERROR: LoadError: no unique matching method found for the specified argument types
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] _which(tt::Type; method_table::Nothing, world::UInt64, raise::Bool)
   @ Base ./reflection.jl:1576
 [3] _which
   @ ./reflection.jl:1563 [inlined]
 [4] which
   @ ./reflection.jl:1602 [inlined]
 [5] which
   @ ./reflection.jl:1593 [inlined]
 [6] #s6#4
   @ ~/Julia/tmp/ConstructionBase/wip.jl:7 [inlined]
 [7] var"#s6#4"(::Any)
   @ Main ./none:0
 [8] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
   @ Core ./boot.jl:600
 [9] top-level scope
   @ ~/Julia/tmp/ConstructionBase/wip.jl:10
in expression starting at /home/tim/Julia/tmp/ConstructionBase/wip.jl:10

For ConstructionBase.jl, calling hasmethod doesn't suffice, because it wants to check if an overload exists.

EDIT: looks like forwarding the new world argument helps, but this obviously affects way more packages than just the GeneratedFunctionStub change...

[vtjnash]

Yes, I was going to point you towards _which and GeneratedFunctionStub (with identity or Returns(bool)), but looks like you found those!

Although _which also throws away edges, so Core.Compiler._findsup is much more algorithmically correct, and then you need to call Meta.lower on the GeneratedFunctionStub result so that you can set those into the results.