diff --git a/Project.toml b/Project.toml index 52e1161..e6edd44 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "CompilerCaching" uuid = "9db33cc3-5358-4881-8759-fa4194144afd" -version = "0.3.1" +version = "0.4.0" [workspace] projects = ["test", "examples"] diff --git a/examples/julia.jl b/examples/julia.jl index 98b499c..0e202c2 100644 --- a/examples/julia.jl +++ b/examples/julia.jl @@ -58,8 +58,8 @@ function getglobal_jljit() end """ - julia_codegen(cache, mi, ci; argtypes=nothing, dump_llvm=false, dump_module=false) - -> (ir_bytes, entry_name, llvm_ir) + julia_codegen(cache, interp, mi, ci; argtypes=nothing, dump_llvm=false, + dump_module=false) -> (ir_bytes, entry_name, llvm_ir) Generate LLVM IR and return serializable intermediate result. Returns a tuple of (LLVM bitcode bytes, entry function name, LLVM IR text). @@ -67,14 +67,15 @@ The `llvm_ir` string is empty unless `dump_llvm` or `dump_module` is set. When `dump_llvm` is true, returns the IR of just the entry function. When `dump_module` is true, returns the IR of the entire module. -Uses `get_codeinfos(ci)` to collect CodeInfos by walking :invoke statements (1.12+) -or cache lookup callback (1.11). When `argtypes` is provided, uses the const-optimized -source for the root CI via `get_codeinfos(ci, argtypes)`. +Uses `get_codeinfos(interp, ci)` to collect CodeInfos by walking :invoke statements +(1.12+, re-inferring through `interp` where cache history left gaps) or the cache +lookup callback (1.11). When `argtypes` is provided, uses the const-optimized source +for the root CI via `get_codeinfos(interp, ci, argtypes)`. This function handles codegen but does not JIT compile - use `julia_jit` for that. """ -function julia_codegen(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance; +function julia_codegen(cache::CacheView, interp::CC.AbstractInterpreter, + mi::Core.MethodInstance, ci::Core.CodeInstance; argtypes::Union{Vector{Any},Nothing}=nothing, dump_llvm::Bool=false, dump_module::Bool=false) @@ -106,7 +107,8 @@ function julia_codegen(cache::CacheView, mi::Core.MethodInstance, # Generate native code @static if VERSION >= v"1.12.0-DEV.1823" cis_vec = Any[] - codeinfos = argtypes !== nothing ? get_codeinfos(ci, argtypes) : get_codeinfos(ci) + codeinfos = argtypes !== nothing ? get_codeinfos(interp, ci, argtypes) : + get_codeinfos(interp, ci) for (ci, src) in codeinfos push!(cis_vec, ci) push!(cis_vec, src) diff --git a/examples/native.jl b/examples/native.jl index 32147e8..6060063 100644 --- a/examples/native.jl +++ b/examples/native.jl @@ -75,9 +75,9 @@ const compilations = Ref(0) # for testing function compile!(cache::CacheView, mi::Core.MethodInstance) # Get a CI through inference + interp = CustomInterpreter(cache) ci = get(cache, mi, nothing) if ci === nothing - interp = CustomInterpreter(cache) ci = CompilerCaching.typeinf!(interp, mi) end @@ -90,7 +90,7 @@ function compile!(cache::CacheView, mi::Core.MethodInstance) # emit code: generate LLVM IR if res.code === nothing - res.code = julia_codegen(cache, mi, ci) + res.code = julia_codegen(cache, interp, mi, ci) end # emit executable: JIT compile to function pointer @@ -102,20 +102,19 @@ function compile!(cache::CacheView, mi::Core.MethodInstance) end function compile!(cache::CacheView, mi::Core.MethodInstance, argtypes::Vector{Any}) + interp = CustomInterpreter(cache) ci = get(cache, mi, nothing) if ci === nothing - interp = CustomInterpreter(cache) ci = CompilerCaching.typeinf!(interp, mi) end # Ensure const-seeded inference has run if CompilerCaching.get_source(ci, argtypes) === nothing - interp = CustomInterpreter(cache) CompilerCaching.typeinf!(cache, interp, mi, argtypes) end # codegen + JIT using const-optimized source - code = julia_codegen(cache, mi, ci; argtypes) + code = julia_codegen(cache, interp, mi, ci; argtypes) return julia_jit(cache, mi, code) end @@ -222,14 +221,14 @@ let CompilerCaching.typeinf!(cache, interp, mi, argtypes) # Generic codegen - (_, _, generic_ir) = julia_codegen(cache, mi, ci; dump_llvm=true) + (_, _, generic_ir) = julia_codegen(cache, interp, mi, ci; dump_llvm=true) println("=== Generic LLVM IR ===") println(generic_ir) @assert contains(generic_ir, "icmp") "Generic IR should have a comparison" @assert contains(generic_ir, "sub i64") "Generic IR should have the sub branch" # Const-seeded codegen (n=3 is a known constant) - (_, _, const_ir) = julia_codegen(cache, mi, ci; argtypes, dump_llvm=true) + (_, _, const_ir) = julia_codegen(cache, interp, mi, ci; argtypes, dump_llvm=true) println("\n=== Const-seeded LLVM IR (n=3) ===") println(const_ir) @assert !contains(const_ir, "icmp") "Const-seeded IR should eliminate the comparison" diff --git a/src/CompilerCaching.jl b/src/CompilerCaching.jl index eb6f4cb..2ad1975 100644 --- a/src/CompilerCaching.jl +++ b/src/CompilerCaching.jl @@ -659,6 +659,11 @@ compiles all callees and stores their source so [`get_codeinfos`](@ref) works. Returns the root `CodeInstance` (or `nothing` if inference failed). Subsequent calls for the same `mi` and world are no-ops — the existing CI is returned. + +The eager callee walk only follows `:invoke` edges that refer to a `CodeInstance`. +Optimized source (in particular source reused from the cache) can also contain +`:invoke` statements targeting a bare `MethodInstance`; those callees are not +compiled here, but are resolved lazily by [`get_codeinfos(interp, ci)`](@ref). """ function typeinf!(interp::CC.AbstractInterpreter, mi::Core.MethodInstance) @static if VERSION >= v"1.12.0-DEV.1434" @@ -1028,27 +1033,60 @@ function get_source(ci::Core.CodeInstance, argtypes::Vector{Any}) end """ - get_codeinfos(ci::CodeInstance) -> Vector{Pair{CodeInstance, CodeInfo}} - -Collect CodeInstance/CodeInfo pairs by walking forward edges from a root CI. - -On Julia 1.12+, walks `:invoke` statements to collect callees transitively. -On Julia 1.11, returns only the root entry. - -Requires that `typeinf!` was called first to populate source for all callees. + get_codeinfos(interp::AbstractInterpreter, ci::CodeInstance) -> + Vector{Pair{CodeInstance, CodeInfo}} + +Collect the `CodeInstance`/`CodeInfo` pairs needed to generate code for `ci`, walking +`:invoke` edges transitively from the root. On Julia 1.12+ the result is closed under +direct-call edges, making it suitable for closed-world code generation +(`jl_emit_native`, which cannot look up missing callees during codegen). + +`interp` is used to repair gaps that cache history can leave behind: + +- `:invoke`/`:invoke_modify` statements whose target is still a bare `MethodInstance` + (inlining's `compileable_specialization` emits those when the compileable + specialization was not cached at optimization time; codegen lowers them to runtime + dispatch) are resolved to a `CodeInstance` — running inference through `interp` if + the cache has none — and the statement is rewritten to target it, in a copy of the + containing source; cached `CodeInfo` is never mutated. Only targets callable through + a native ABI (concrete signature, fully-resolved sparams) are resolved; anything + else (e.g. `@nospecialize`-widened compileable signatures) deliberately keeps its + runtime-dispatch fallback semantics. +- Referenced `CodeInstance`s that lack stored source (e.g. cached by an earlier + session) are re-inferred. + +`interp` must match the cache owner and world that produced `ci` (typically the +interpreter previously passed to [`typeinf!`](@ref)). + +On Julia 1.11, code generation resolves callees through a lookup callback instead, so +only the root entry is returned. """ -function get_codeinfos(ci::Core.CodeInstance) +get_codeinfos(interp::CC.AbstractInterpreter, ci::Core.CodeInstance) = + collect_codeinfos(interp, ci, nothing) + +function collect_codeinfos(interp::CC.AbstractInterpreter, + root::Core.CodeInstance, root_src::Union{Core.CodeInfo, Nothing}) codeinfos = Pair{Core.CodeInstance, Core.CodeInfo}[] @static if VERSION >= v"1.12-" visited = IdSet{Core.CodeInstance}() - workqueue = Core.CodeInstance[ci] + workqueue = Core.CodeInstance[root] while !isempty(workqueue) callee_ci = pop!(workqueue) callee_ci in visited && continue push!(visited, callee_ci) - src = get_source(callee_ci) - @assert src !== nothing "CodeInstance for $(CC.get_ci_mi(callee_ci)) has no source - ensure typeinf! was called" + src = callee_ci === root && root_src !== nothing ? root_src : + get_source(callee_ci) + if src === nothing + # a referenced CI may lack stored source (e.g. it was cached by an + # earlier session whose sources were dropped); re-establish it + typeinf!(interp, CC.get_ci_mi(callee_ci)) + src = get_source(callee_ci) + # if inference cannot provide source either, leave the call site to + # codegen's runtime-dispatch fallback + src === nothing && continue + end + src = resolve_invoke_targets(interp, src) push!(codeinfos, callee_ci => src) for stmt in src.code @@ -1064,40 +1102,78 @@ function get_codeinfos(ci::Core.CodeInstance) end end else - src = get_source(ci) - src !== nothing && push!(codeinfos, ci => src) + src = root_src === nothing ? get_source(root) : root_src + src !== nothing && push!(codeinfos, root => src) end return codeinfos end -""" - get_codeinfos(ci::CodeInstance, argtypes::Vector{Any}) -> Vector{Pair{CodeInstance, CodeInfo}} +@static if VERSION >= v"1.12-" +# Mirror of nightly's `Compiler.has_valid_abi_sparams`: specializations with +# incomplete sparams (TypeVar) or SimpleVector/Vararg sparams cannot be called +# through a native specsig ABI — codegen's `needsparams` path emits `jl_invoke` +# even for CodeInstance operands, so rewriting such targets is useless. +function has_valid_abi_sparams(mi::Core.MethodInstance) + for sp in mi.sparam_vals + if sp isa TypeVar || sp isa Core.SimpleVector || CC.isvarargtype(sp) + return false + end + end + return true +end -Collect CodeInstance/CodeInfo pairs using the const-optimized source for the root CI -and generic source for all callees. +# Rewrite `:invoke`/`:invoke_modify` statements whose target is still a bare +# `MethodInstance` to target a `CodeInstance` instead, inferring one when the cache +# has none. Codegen can only emit a direct call for a `CodeInstance` operand; a +# `MethodInstance` operand unconditionally lowers to runtime dispatch. Returns `src` +# unchanged when there is nothing to rewrite, or a rewritten copy (cached source is +# never mutated). +function resolve_invoke_targets(interp::CC.AbstractInterpreter, src::Core.CodeInfo) + resolved = src + for pc in eachindex(resolved.code) + stmt = resolved.code[pc] + rhs = stmt isa Expr && stmt.head === :(=) ? stmt.args[2] : stmt + rhs isa Expr && (rhs.head === :invoke || rhs.head === :invoke_modify) || continue + + callee_mi = rhs.args[1] + callee_mi isa Core.MethodInstance || continue + # Only specializations with a fully-concrete signature can be invoked directly + # through a native ABI; anything else (e.g. `@nospecialize`-widened compileable + # signatures) keeps its runtime-dispatch fallback semantics. + callee_mi.def isa Method && isdispatchtuple(callee_mi.specTypes) && + has_valid_abi_sparams(callee_mi) || continue + + callee_ci = typeinf!(interp, callee_mi) + callee_ci === nothing && continue + # only rewrite when the callee's source is available, so the returned + # collection stays closed under direct-call edges + get_source(callee_ci) === nothing && continue + + if resolved === src + resolved = copy(src) + stmt = resolved.code[pc] + rhs = stmt isa Expr && stmt.head === :(=) ? stmt.args[2] : stmt + end + new_rhs = Expr(rhs.head, callee_ci, rhs.args[2:end]...) + resolved.code[pc] = stmt === rhs ? new_rhs : Expr(:(=), stmt.args[1], new_rhs) + end + return resolved +end +end + +""" + get_codeinfos(interp::AbstractInterpreter, ci::CodeInstance, argtypes::Vector{Any}) -> + Vector{Pair{CodeInstance, CodeInfo}} -Delegates to `get_codeinfos(ci)` for the full callee walk, then swaps the root entry's -source with the const-optimized version from `get_source(ci, argtypes)`. Extra callees -from the generic walk are harmless (compiled but uncalled), and missing callees get -runtime dispatch stubs from `jl_emit_native`. +Const-specialized variant of [`get_codeinfos(interp, ci)`](@ref): the callee walk is +seeded with the const-optimized source stored for `argtypes` (see +[`typeinf!(cache, interp, mi, argtypes)`](@ref)), so callees reachable only from the +const-optimized code are included as well. -Falls back to `get_codeinfos(ci)` if no const entry exists for the given argtypes. +Falls back to the generic source if no const entry exists for the given argtypes. """ -function get_codeinfos(ci::Core.CodeInstance, argtypes::Vector{Any}) - const_src = get_source(ci, argtypes) - if const_src === nothing - return get_codeinfos(ci) - end - codeinfos = get_codeinfos(ci) - # Swap root entry's source with const-optimized version - idx = findfirst(p -> p.first === ci, codeinfos) - if idx !== nothing - codeinfos[idx] = ci => const_src - else - pushfirst!(codeinfos, ci => const_src) - end - return codeinfos -end +get_codeinfos(interp::CC.AbstractInterpreter, ci::Core.CodeInstance, argtypes::Vector{Any}) = + collect_codeinfos(interp, ci, get_source(ci, argtypes)) end # @static if VERSION >= v"1.11" diff --git a/test/runtests.jl b/test/runtests.jl index fe776ae..c6ab71d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -435,6 +435,123 @@ end # with healthy const-prop, the `:b` branch (and its Ptr{Int32} store) is elided @test !any(stmt -> occursin("Int32", string(stmt)), src.code) end + +@testset "unresolved invoke targets" begin + # Inlining's `compileable_specialization` leaves an `:invoke` target as a bare + # MethodInstance when the compileable specialization is not cached at optimization + # time; codegen lowers such statements to runtime dispatch. Cached optimized + # source can therefore reference callees that a plain CodeInstance-edge walk + # cannot follow (JuliaGPU/Metal.jl checked-conversion overlay failures on 1.12+). + # `get_codeinfos(interp, ci)` must resolve concrete targets to CodeInstances. + mod = @eval module $(gensym()) + using Base.Experimental: @MethodTable, @overlay + @MethodTable method_table + function overlay_child end + @overlay method_table @noinline overlay_child(x::UInt32, p::Ptr{Int32}) = + unsafe_store!(p, x % Int32) + + @noinline child(x::UInt32, p::Ptr{Int32}) = unsafe_store!(p, x % Int32) + @noinline nospec_child(p::Ptr{Int32}, @nospecialize(x)) = + unsafe_store!(p, Int32(2)) + function kernel(x::UInt32, p::Ptr{Int32}) + child(x, p) + nospec_child(p, x) + return + end + end + + get_ci_mi = Core.Compiler.get_ci_mi + function find_invoke(src::Core.CodeInfo, pred) + findfirst(eachindex(src.code)) do pc + stmt = src.code[pc] + rhs = stmt isa Expr && stmt.head === :(=) ? stmt.args[2] : stmt + rhs isa Expr && rhs.head === :invoke && pred(rhs.args[1]) + end + end + invoke_target(src::Core.CodeInfo, pc::Int) = begin + stmt = src.code[pc] + rhs = stmt isa Expr && stmt.head === :(=) ? stmt.args[2] : stmt + rhs.args[1] + end + + world = Base.get_world_counter() + cache = CacheView{TestResults}(:UnresolvedInvokeTest, world) + interp = TestInterpreter(cache.world, cache, InfCacheT()) + + mi = method_instance(mod.kernel, (UInt32, Ptr{Int32}); world) + root = typeinf!(interp, mi) + @test root isa Core.CodeInstance + + overlay_mi = method_instance(mod.overlay_child, (UInt32, Ptr{Int32}); + world, method_table=mod.method_table) + @test overlay_mi isa Core.MethodInstance + + # sources that need no rewriting are returned as-is, not copied (only checkable + # when the stored source is an uncompressed CodeInfo) + stored = @atomic :monotonic root.inferred + if stored isa Core.CodeInfo + @test only(src for (ci, src) in get_codeinfos(interp, root) if ci === root) === + stored + end + + # reconstruct the problematic cache state: retarget the `child` invoke (a + # CodeInstance edge from fresh inference) to the never-inferred overlay MI + doctored = copy(get_source(root)) + child_pc = find_invoke(doctored, op -> op isa Core.CodeInstance && + get_ci_mi(op).def in methods(mod.child)) + @test child_pc !== nothing + stmt = doctored.code[child_pc] + rhs = stmt isa Expr && stmt.head === :(=) ? stmt.args[2] : stmt + new_rhs = Expr(:invoke, overlay_mi, rhs.args[2:end]...) + doctored.code[child_pc] = stmt === rhs ? new_rhs : Expr(:(=), stmt.args[1], new_rhs) + @atomic root.inferred = doctored + + # the walk infers the unresolved callee, includes it, and rewrites the operand + pairs = get_codeinfos(interp, root) + overlay_idx = findfirst(p -> get_ci_mi(p.first) === overlay_mi, pairs) + @test overlay_idx !== nothing + root_src = only(src for (ci, src) in pairs if ci === root) + @test invoke_target(root_src, child_pc) === pairs[overlay_idx].first + # copy-on-write: the cached source still carries the MethodInstance operand + @test root_src !== doctored + @test invoke_target(doctored, child_pc) === overlay_mi + + # targets with non-concrete signatures (@nospecialize-widened compileable + # signatures) keep their runtime-dispatch fallback + nospec_pc = find_invoke(doctored, op -> begin + op_mi = op isa Core.CodeInstance ? get_ci_mi(op) : op + op_mi isa Core.MethodInstance && op_mi.def in methods(mod.nospec_child) + end) + @test nospec_pc !== nothing + if invoke_target(doctored, nospec_pc) isa Core.MethodInstance + @test invoke_target(root_src, nospec_pc) === invoke_target(doctored, nospec_pc) + @test !any(p -> get_ci_mi(p.first).def in methods(mod.nospec_child), pairs) + end + + # const-specialized collection receives the same treatment, seeded from the + # const-optimized root source + const_argtypes = Any[Core.Compiler.Const(mod.kernel), Core.Compiler.Const(UInt32(42)), + Ptr{Int32}] + typeinf!(cache, interp, mi, const_argtypes) + const_src = get_source(root, const_argtypes) + @test const_src isa Core.CodeInfo + const_pc = find_invoke(const_src, op -> op isa Core.CodeInstance && + get_ci_mi(op).def in methods(mod.child)) + @test const_pc !== nothing + stmt = const_src.code[const_pc] + rhs = stmt isa Expr && stmt.head === :(=) ? stmt.args[2] : stmt + new_rhs = Expr(:invoke, overlay_mi, rhs.args[2:end]...) + const_src.code[const_pc] = stmt === rhs ? new_rhs : Expr(:(=), stmt.args[1], new_rhs) + + const_pairs = get_codeinfos(interp, root, const_argtypes) + @test any(p -> get_ci_mi(p.first) === overlay_mi, const_pairs) + const_root_src = only(src for (ci, src) in const_pairs if ci === root) + @test invoke_target(const_root_src, const_pc) isa Core.CodeInstance + @test get_ci_mi(invoke_target(const_root_src, const_pc)) === overlay_mi + # the stored const-prop entry itself is left unmutated + @test const_root_src !== const_src + @test invoke_target(const_src, const_pc) === overlay_mi +end end #==============================================================================#