Adapt handling of world ages to upstream changes (#403)

Author: maleadt

src/cache.jl

@@ -38,16 +38,98 @@ Normally, you shouldn't have to use this function, as it's used by `FunctionSpec
 """
 get_world
 
-# generate functions currently do not know which world they are invoked for, so we fall
-# back to using the current world. this may be wrong when the generator is invoked in a
-# different world (TODO: when does this happen?)
-#
-# XXX: this should be fixed by JuliaLang/julia#48611
+if VERSION >= v"1.10.0-DEV.873"
 
-function get_world_generator(self, ::Type{Type{ft}}, ::Type{Type{tt}}) where {ft, tt}
+# on 1.10 (JuliaLang/julia#48611) the generated function knows which world it was invoked in
+
+function _generated_ex(world, source, ex)
+    stub = Core.GeneratedFunctionStub(identity, Core.svec(:get_world, :ft, :tt), Core.svec())
+    stub(world, source, ex)
+end
+
+function get_world_generator(world::UInt, source, self, ft::Type, tt::Type)
     @nospecialize
+    @assert Core.Compiler.isType(ft) && Core.Compiler.isType(tt)
+    ft = ft.parameters[1]
+    tt = tt.parameters[1]
 
     # look up the method
+    method_error = :(throw(MethodError(ft, tt, $world)))
+    Base.isdispatchtuple(tt) || return _generated_ex(world, source, :(error("$tt is not a dispatch tuple")))
+    sig = Tuple{ft, tt.parameters...}
+    min_world = Ref{UInt}(typemin(UInt))
+    max_world = Ref{UInt}(typemax(UInt))
+    has_ambig = Ptr{Int32}(C_NULL)  # don't care about ambiguous results
+    mthds = if VERSION >= v"1.7.0-DEV.1297"
+        Base._methods_by_ftype(sig, #=mt=# nothing, #=lim=# -1,
+                               world, #=ambig=# false,
+                               min_world, max_world, has_ambig)
+        # XXX: use the correct method table to support overlaying kernels
+    else
+        Base._methods_by_ftype(sig, #=lim=# -1,
+                               world, #=ambig=# false,
+                               min_world, max_world, has_ambig)
+    end
+    mthds === nothing && return _generated_ex(world, source, method_error)
+    length(mthds) == 1 || return _generated_ex(world, source, method_error)
+
+    # look up the method and code instance
+    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)                # see note below
+    empty!(new_ci.ssaflags)
+    new_ci.ssavaluetypes = 0
+    new_ci.min_world = min_world[]
+    new_ci.max_world = max_world[]
+    new_ci.edges = MethodInstance[mi]
+    # XXX: setting this edge does not give us proper method invalidation, see
+    #      JuliaLang/julia#34962 which demonstrates we also need to "call" the kernel.
+    #      invoking `code_llvm` also does the necessary codegen, as does calling the
+    #      underlying C methods -- which GPUCompiler does, so everything Just Works.
+
+    # prepare the slots
+    new_ci.slotnames = Symbol[Symbol("#self#"), :ft, :tt]
+    new_ci.slotflags = UInt8[0x00 for i = 1:3]
+
+    # return the world
+    push!(new_ci.code, ReturnNode(world))
+    push!(new_ci.ssaflags, 0x00)   # Julia's native compilation pipeline (and its verifier) expects `ssaflags` to be the same length as `code`
+    push!(new_ci.codelocs, 1)   # see note below
+    new_ci.ssavaluetypes += 1
+
+    # NOTE: we keep the first entry of the original linetable, and use it for location info
+    #       on the call to check_cache. we can't not have a codeloc (using 0 causes
+    #       corruption of the back trace), and reusing the target function's info
+    #       has as advantage that we see the name of the kernel in the backtraces.
+
+    return new_ci
+end
+
+@eval function get_world(ft, tt)
+    $(Expr(:meta, :generated_only))
+    $(Expr(:meta, :generated, get_world_generator))
+end
+
+else
+
+# on older versions of Julia we fall back to looking up the current world. this may be wrong
+# when the generator is invoked in a different world (TODO: when does this happen?)
+
+function get_world_generator(self, ft::Type, tt::Type)
+    @nospecialize
+    @assert Core.Compiler.isType(ft) && Core.Compiler.isType(tt)
+    ft = ft.parameters[1]
+    tt = tt.parameters[1]
+
+    # look up the method
+    method_error = :(throw(MethodError(ft, tt)))
+    Base.isdispatchtuple(tt) || return(:(error("$tt is not a dispatch tuple")))
     sig = Tuple{ft, tt.parameters...}
     min_world = Ref{UInt}(typemin(UInt))
     max_world = Ref{UInt}(typemax(UInt))
@@ -63,11 +145,7 @@ function get_world_generator(self, ::Type{Type{ft}}, ::Type{Type{tt}}) where {ft
                                min_world, max_world, has_ambig)
     end
     # XXX: using world=-1 is wrong, but the current world isn't exposed to this generator
-
-    # check the validity of the method matches
-    method_error = :(throw(MethodError(ft, tt)))
     mthds === nothing && return method_error
-    Base.isdispatchtuple(tt) || return(:(error("$tt is not a dispatch tuple")))
     length(mthds) == 1 || return method_error
 
     # look up the method and code instance
@@ -127,6 +205,8 @@ end
                 true)))
 end
 
+end
+
 const cache_lock = ReentrantLock()
 
 """