typeinfer: add separate compile workqueue for native / call_latest code

Author: topolarity

Compiler/src/typeinfer.jl

@@ -1378,6 +1378,12 @@ end
 # Method and return a compilable MethodInstance for the call, if
 # it will be runtime-dispatched to exactly that MethodInstance
 function compileable_specialization_for_call(interp::AbstractInterpreter, @nospecialize(argtype))
+    mt = ccall(:jl_method_table_for, Any, (Any,), argtype)
+    if mt === nothing
+        # this would require scanning all method tables, so give up instead
+        return nothing
+    end
+
     matches = findall(argtype, method_table(interp); limit = 1)
     matches === nothing && return nothing
     length(matches.matches) == 0 && return nothing
@@ -1392,41 +1398,64 @@ function compileable_specialization_for_call(interp::AbstractInterpreter, @nospe
     else
         mi = specialize_method(match.method, compileable_atype, match.sparams)
     end
+
     return mi
 end
 
+const QueueItems = Union{CodeInstance,MethodInstance,SimpleVector}
+
+struct CompilationQueue
+    tocompile::Vector{QueueItems}
+    inspected::IdSet{QueueItems}
+    interp::Union{AbstractInterpreter,Nothing}
+
+    CompilationQueue(;
+        interp::Union{AbstractInterpreter,Nothing}
+    ) = new(QueueItems[], IdSet{QueueItems}(), interp)
+
+    CompilationQueue(queue::CompilationQueue;
+        interp::Union{AbstractInterpreter,Nothing}
+    ) = new(empty!(queue.tocompile), empty!(queue.inspected), interp)
+end
+
+Base.push!(queue::CompilationQueue, item) = push!(queue.tocompile, item)
+Base.append!(queue::CompilationQueue, items) = append!(queue.tocompile, items)
+Base.pop!(queue::CompilationQueue) = pop!(queue.tocompile)
+Base.empty!(queue::CompilationQueue) = (empty!(queue.tocompile); empty!(queue.inspected))
+markinspected!(queue::CompilationQueue, item) = push!(queue.inspected, item)
+isinspected(queue::CompilationQueue, item) = item in queue.inspected
+Base.isempty(queue::CompilationQueue) = isempty(queue.tocompile)
+
 # collect a list of all code that is needed along with CodeInstance to codegen it fully
-function collectinvokes!(wq::Vector{CodeInstance}, ci::CodeInfo, sptypes::Vector{VarState})
+function collectinvokes!(workqueue::CompilationQueue, ci::CodeInfo, sptypes::Vector{VarState};
+                         invokelatest_queue::Union{CompilationQueue,Nothing} = nothing)
     src = ci.code
     for i = 1:length(src)
         stmt = src[i]
         isexpr(stmt, :(=)) && (stmt = stmt.args[2])
         if isexpr(stmt, :invoke) || isexpr(stmt, :invoke_modify)
             edge = stmt.args[1]
-            edge isa CodeInstance && isdefined(edge, :inferred) && push!(wq, edge)
+            edge isa CodeInstance && isdefined(edge, :inferred) && push!(workqueue, edge)
         end
 
         if isexpr(stmt, :call)
             farg = stmt.args[1]
             !applicable(argextype, farg, ci, sptypes) && continue # TODO: Why is this failing during bootstrap
             ftyp = widenconst(argextype(farg, ci, sptypes))
             if ftyp <: Builtin
-                # TODO: Make interp elsewhere
-                interp = NativeInterpreter(Base.get_world_counter())
                 if Core.finalizer isa ftyp && length(stmt.args) == 3
                     finalizer = argextype(stmt.args[2], ci, sptypes)
                     obj = argextype(stmt.args[3], ci, sptypes)
                     atype = argtypes_to_type(Any[finalizer, obj])
 
-                    mi = compileable_specialization_for_call(interp, atype)
-                    mi === nothing && continue
+                    invokelatest_queue === nothing && continue
+                    let workqueue = invokelatest_queue
+                        # make a best-effort attempt to enqueue the relevant code for the finalizer
+                        mi = compileable_specialization_for_call(workqueue.interp, atype)
+                        mi === nothing && continue
 
-                    if mi.def.primary_world <= Base.get_world_counter() <= mi.def.deleted_world
-                        ci = typeinf_ext(interp, mi, SOURCE_MODE_GET_SOURCE)
-                        # TODO: separate workqueue for NativeInterpreter
-                        ci isa CodeInstance && push!(wq, ci)
+                        push!(workqueue, mi)
                     end
-                    # push!(wq, mi)
                 end
             end
         end
@@ -1439,41 +1468,40 @@ function add_codeinsts_to_jit!(interp::AbstractInterpreter, ci, source_mode::UIn
     ci isa CodeInstance && !ci_has_invoke(ci) || return ci
     codegen = codegen_cache(interp)
     codegen === nothing && return ci
-    inspected = IdSet{CodeInstance}()
-    tocompile = Vector{CodeInstance}()
-    push!(tocompile, ci)
-    while !isempty(tocompile)
+    workqueue = CompilationQueue(; interp)
+    push!(workqueue, ci)
+    while !isempty(workqueue)
         # ci_has_real_invoke(ci) && return ci # optimization: cease looping if ci happens to get compiled (not just jl_fptr_wait_for_compiled, but fully jl_is_compiled_codeinst)
-        callee = pop!(tocompile)
+        callee = pop!(workqueue)
         ci_has_invoke(callee) && continue
-        callee in inspected && continue
+        isinspected(workqueue, callee) && continue
         src = get(codegen, callee, nothing)
         if !isa(src, CodeInfo)
             src = @atomic :monotonic callee.inferred
             if isa(src, String)
                 src = _uncompressed_ir(callee, src)
             end
             if !isa(src, CodeInfo)
-                newcallee = typeinf_ext(interp, callee.def, source_mode) # always SOURCE_MODE_ABI
+                newcallee = typeinf_ext(workqueue.interp, callee.def, source_mode) # always SOURCE_MODE_ABI
                 if newcallee isa CodeInstance
                     callee === ci && (ci = newcallee) # ci stopped meeting the requirements after typeinf_ext last checked, try again with newcallee
-                    push!(tocompile, newcallee)
+                    push!(workqueue, newcallee)
                 end
                 if newcallee !== callee
-                    push!(inspected, callee)
+                    markinspected!(workqueue, callee)
                 end
                 continue
             end
         end
-        push!(inspected, callee)
+        markinspected!(workqueue, callee)
         mi = get_ci_mi(callee)
         sptypes = sptypes_from_meth_instance(mi)
-        collectinvokes!(tocompile, src, sptypes)
+        collectinvokes!(workqueue, src, sptypes)
         if iszero(ccall(:jl_mi_cache_has_ci, Cint, (Any, Any), mi, callee))
-            cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, get_inference_world(interp))::CodeInstance
+            cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, get_inference_world(workqueue.interp))::CodeInstance
             if cached === callee
                 # make sure callee is gc-rooted and cached, as required by jl_add_codeinst_to_jit
-                code_cache(interp)[mi] = callee
+                code_cache(workqueue.interp)[mi] = callee
             else
                 # use an existing CI from the cache, if there is available one that is compatible
                 callee === ci && (ci = cached)
@@ -1497,57 +1525,45 @@ function typeinf_ext_toplevel(mi::MethodInstance, world::UInt, source_mode::UInt
     return typeinf_ext_toplevel(interp, mi, source_mode)
 end
 
-# This is a bridge for the C code calling `jl_typeinf_func()` on set of Method matches
-# The trim_mode can be any of:
-const TRIM_NO = 0
-const TRIM_SAFE = 1
-const TRIM_UNSAFE = 2
-const TRIM_UNSAFE_WARN = 3
-function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_mode::Int)
-    inspected = IdSet{CodeInstance}()
-    tocompile = Vector{CodeInstance}()
-    codeinfos = []
-    # first compute the ABIs of everything
-    latest = true # whether this_world == world_counter()
-    for this_world in reverse(sort!(worlds))
-        interp = NativeInterpreter(
-            this_world;
-            inf_params = InferenceParams(; force_enable_inference = trim_mode != TRIM_NO)
-        )
-        for i = 1:length(methods)
-            # each item in this list is either a MethodInstance indicating something
-            # to compile, or an svec(rettype, sig) describing a C-callable alias to create.
-            item = methods[i]
-            if item isa MethodInstance
-                # if this method is generally visible to the current compilation world,
-                # and this is either the primary world, or not applicable in the primary world
-                # then we want to compile and emit this
-                if item.def.primary_world <= this_world <= item.def.deleted_world
-                    ci = typeinf_ext(interp, item, SOURCE_MODE_GET_SOURCE)
-                    ci isa CodeInstance && push!(tocompile, ci)
-                end
-            elseif item isa SimpleVector && latest
-                (rt::Type, sig::Type) = item
-                # make a best-effort attempt to enqueue the relevant code for the ccallable
-                ptr = ccall(:jl_get_specialization1,
-                            #= MethodInstance =# Ptr{Cvoid}, (Any, Csize_t, Cint),
-                            sig, this_world, #= mt_cache =# 0)
-                if ptr !== C_NULL
-                    mi = unsafe_pointer_to_objref(ptr)::MethodInstance
-                    ci = typeinf_ext(interp, mi, SOURCE_MODE_GET_SOURCE)
-                    ci isa CodeInstance && push!(tocompile, ci)
-                end
-                # additionally enqueue the ccallable entrypoint / adapter, which implicitly
-                # invokes the above ci
-                push!(codeinfos, item)
+function compile!(codeinfos::Vector{Any}, workqueue::CompilationQueue;
+    invokelatest_queue::Union{CompilationQueue,Nothing} = nothing,
+)
+    interp = workqueue.interp
+    world = get_inference_world(interp)
+    while !isempty(workqueue)
+        item = pop!(workqueue)
+        # each item in this list is either a MethodInstance indicating something
+        # to compile, or an svec(rettype, sig) describing a C-callable alias to create.
+        if item isa MethodInstance
+            isinspected(workqueue, item) && continue
+            # if this method is generally visible to the current compilation world,
+            # and this is either the primary world, or not applicable in the primary world
+            # then we want to compile and emit this
+            if item.def.primary_world <= world <= item.def.deleted_world
+                ci = typeinf_ext(interp, item, SOURCE_MODE_GET_SOURCE)
+                ci isa CodeInstance && push!(workqueue, ci)
             end
-        end
-        while !isempty(tocompile)
-            callee = pop!(tocompile)
-            callee in inspected && continue
-            # now make sure everything has source code, if desired
+            markinspected!(workqueue, item)
+        elseif item isa SimpleVector
+            invokelatest_queue === nothing && continue
+            (rt::Type, sig::Type) = item
+            # make a best-effort attempt to enqueue the relevant code for the ccallable
+            ptr = ccall(:jl_get_specialization1,
+                        #= MethodInstance =# Ptr{Cvoid}, (Any, Csize_t, Cint),
+                        sig, world, #= mt_cache =# 0)
+            if ptr !== C_NULL
+                mi = unsafe_pointer_to_objref(ptr)::MethodInstance
+                ci = typeinf_ext(interp, mi, SOURCE_MODE_GET_SOURCE)
+                ci isa CodeInstance && push!(invokelatest_queue, ci)
+            end
+            # additionally enqueue the ccallable entrypoint / adapter, which implicitly
+            # invokes the above ci
+            push!(codeinfos, item)
+        elseif item isa CodeInstance
+            callee = item
+            isinspected(workqueue, callee) && continue
             mi = get_ci_mi(callee)
-            def = mi.def
+            # now make sure everything has source code, if desired
             if use_const_api(callee)
                 src = codeinfo_for_const(interp, mi, callee.rettype_const)
             else
@@ -1556,21 +1572,21 @@ function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_m
                     newcallee = typeinf_ext(interp, mi, SOURCE_MODE_GET_SOURCE)
                     if newcallee isa CodeInstance
                         @assert use_const_api(newcallee) || haskey(interp.codegen, newcallee)
-                        push!(tocompile, newcallee)
+                        push!(workqueue, newcallee)
                     end
                     if newcallee !== callee
-                        push!(inspected, callee)
+                        markinspected!(workqueue, callee)
                     end
                     continue
                 end
             end
-            push!(inspected, callee)
+            markinspected!(workqueue, callee)
             if src isa CodeInfo
                 sptypes = sptypes_from_meth_instance(mi)
-                collectinvokes!(tocompile, src, sptypes)
+                collectinvokes!(workqueue, src, sptypes; invokelatest_queue)
                 # try to reuse an existing CodeInstance from before to avoid making duplicates in the cache
                 if iszero(ccall(:jl_mi_cache_has_ci, Cint, (Any, Any), mi, callee))
-                    cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, this_world)::CodeInstance
+                    cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, world)::CodeInstance
                     if cached === callee
                         code_cache(interp)[mi] = callee
                     else
@@ -1581,9 +1597,45 @@ function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_m
                 push!(codeinfos, callee)
                 push!(codeinfos, src)
             end
+        else @assert false "unexpected item in queue" end
+    end
+    return codeinfos
+end
+
+# This is a bridge for the C code calling `jl_typeinf_func()` on set of Method matches
+# The trim_mode can be any of:
+const TRIM_NO = 0
+const TRIM_SAFE = 1
+const TRIM_UNSAFE = 2
+const TRIM_UNSAFE_WARN = 3
+function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_mode::Int)
+    inf_params = InferenceParams(; force_enable_inference = trim_mode != TRIM_NO)
+    invokelatest_queue = CompilationQueue(;
+        interp = NativeInterpreter(get_world_counter(); inf_params)
+    )
+    codeinfos = []
+    is_latest_world = true # whether this_world == world_counter()
+    workqueue = CompilationQueue(; interp = nothing)
+    for (i, this_world) in enumerate(sort!(worlds))
+        workqueue = CompilationQueue(workqueue;
+            interp = NativeInterpreter(this_world; inf_params)
+        )
+
+        append!(workqueue, methods)
+
+        is_latest_world = (i == length(worlds))
+        if is_latest_world
+            # Provide the `invokelatest` queue so that we trigger "best-effort" code generation
+            # for, e.g., finalizers and cfunction.
+            #
+            # The queue is intentionally aliased, to handle e.g. a `finalizer` calling `Core.finalizer`
+            # (it will enqueue into itself and immediately drain)
+            compile!(codeinfos, workqueue; invokelatest_queue = workqueue)
+        else
+            compile!(codeinfos, workqueue)
         end
-        latest = false
     end
+
     if trim_mode != TRIM_NO && trim_mode != TRIM_UNSAFE
         verify_typeinf_trim(codeinfos, trim_mode == TRIM_UNSAFE_WARN)
     end