diff --git a/Project.toml b/Project.toml index b1ef7807..eca10a9f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GPUCompiler" uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "2.1.1" +version = "2.2.0" authors = ["Tim Besard "] [workspace] @@ -36,8 +36,8 @@ CompilerCaching = "0.4" ExprTools = "0.1" Highlights = "0.6" InteractiveUtils = "1" -LLVM = "9.9" -LLVMDowngrader_jll = "0.8" +LLVM = "9.11" +LLVMDowngrader_jll = "0.8.3" Libdl = "1" Logging = "1" NVPTX_LLVM_Backend_jll = "22" diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl index d1ef48e2..98bb256d 100644 --- a/src/GPUCompiler.jl +++ b/src/GPUCompiler.jl @@ -54,6 +54,7 @@ include("mangling.jl") # compiler interface and implementations include("interface.jl") +include("relocation.jl") include("error.jl") include("native.jl") include("ptx.jl") @@ -84,13 +85,10 @@ include("precompile.jl") function __init__() STDERR_HAS_COLOR[] = get(stderr, :color, false) + empty!(session_results_cache) @static if !HAS_INTEGRATED_CACHE - # session-local results keyed by CodeInstance; entries serialized during - # GPUCompiler's own precompilation can never be valid in a later session - empty!(legacy_job_results) - # ditto for the in-process CodeCaches: CIs deposited by our own precompile - # workload carry world ages from the precompilation process + # CodeInstances created by GPUCompiler's precompile workload are process-local. empty!(GLOBAL_CI_CACHES) end diff --git a/src/deprecated.jl b/src/deprecated.jl index d3c649eb..e1b80f77 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -229,35 +229,6 @@ function CC.findsup(@nospecialize(sig::Type), table::StackedMethodTable) end -## 1.10 `cached_results` -# -# Session-local storage for the per-job results structs; on 1.11+ these live on the -# `CodeInstance`s of Julia's integrated cache instead (see `interface.jl`). Keep the -# same identity here by associating results with a foreign `CodeInstance`: unrelated -# world-age advances can reuse a still-valid CI, while invalidation makes the lookup -# resolve to a new CI and therefore a new results struct. -struct LegacyJobResultEntry - config::CompilerConfig - value::Any -end - -const legacy_job_results = IdDict{CodeInstance,Vector{LegacyJobResultEntry}}() -const job_results_lock = ReentrantLock() - -function job_results(::Type{V}, ci::CodeInstance, config::CompilerConfig) where {V} - Base.@lock job_results_lock begin - entries = get!(legacy_job_results, ci) do - LegacyJobResultEntry[] - end - for entry in entries - entry.config === config && entry.value isa V && return entry.value::V - end - v = V() - push!(entries, LegacyJobResultEntry(config, v)) - return v - end -end - function job_code_instance(@nospecialize(job::CompilerJob)) cache = WorldView(get_code_cache(job), job.world, job.world) CC.get(cache, job.source, nothing) @@ -268,18 +239,9 @@ end function cached_results(::Type{V}, job::CompilerJob) where {V} ci = job_code_instance(job) ci === nothing && return nothing - return job_results(V, ci, job.config) + return session_results(V, ci, job.config) end - -## 1.10 session-dependent results -# -# Nothing to wipe: `legacy_job_results` never persists meaningfully across sessions (its -# CodeInstance keys are session-specific, and our own image's entries are cleared in -# `__init__`; entries written by a downstream package's workload don't make it into that -# package's image at all, cf. cross-image mutation loss). -mark_session_dependent!(@nospecialize(job::CompilerJob)) = nothing - end # !HAS_INTEGRATED_CACHE diff --git a/src/driver.jl b/src/driver.jl index 38581053..43168a75 100644 --- a/src/driver.jl +++ b/src/driver.jl @@ -50,6 +50,9 @@ export compile Compile a `job` to one of the following formats as specified by the `target` argument: `:llvm` for LLVM IR, `:asm` for assembly, or `:obj` for object code. + +The default [`relocation_lowering`](@ref) strategy resolves Julia-value relocations in the +`:llvm` result. Other strategies retain relocation metadata for their loader. """ function compile(target::Symbol, @nospecialize(job::CompilerJob)) if compile_hook[] !== nothing @@ -58,7 +61,8 @@ function compile(target::Symbol, @nospecialize(job::CompilerJob)) return compile_unhooked(target, job) end -function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob)) +function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob); + resolve_relocations::Bool=true) if context(; throw_error=false) === nothing error("No active LLVM context. Use `JuliaContext()` do-block syntax to create one.") end @@ -73,7 +77,7 @@ function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob)) ## LLVM IR - ir, ir_meta = emit_llvm(job) + ir, ir_meta = emit_llvm(job; resolve_relocations) if output == :llvm if job.config.strip @@ -93,7 +97,7 @@ function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob)) else error("Unknown assembly format $output") end - asm, asm_meta = emit_asm(job, ir, format) + asm, asm_meta = emit_asm(job, ir, ir_meta.relocations, format) if output == :asm || output == :obj return asm, (; asm_meta..., ir_meta..., ir) @@ -169,7 +173,8 @@ end const __llvm_initialized = Ref(false) -@locked function emit_llvm(@nospecialize(job::CompilerJob)) +@locked function emit_llvm(@nospecialize(job::CompilerJob); + resolve_relocations::Bool=true) if !__llvm_initialized[] InitializeAllTargets() InitializeAllTargetInfos() @@ -180,7 +185,7 @@ const __llvm_initialized = Ref(false) end @tracepoint "IR generation" begin - ir, compiled, gv_to_value = irgen(job) + ir, compiled, relocations = irgen(job) if job.config.entry_abi === :specfunc entry_fn = compiled[job.source].specfunc else @@ -245,11 +250,9 @@ const __llvm_initialized = Ref(false) dyn_ir, dyn_meta = @invokelatest deferred_codegen(dyn_job, job) dyn_entry_fn = LLVM.name(dyn_meta.entry) merge!(compiled, dyn_meta.compiled) - if haskey(dyn_meta, :gv_to_value) - merge!(gv_to_value, dyn_meta.gv_to_value) - end @assert context(dyn_ir) == context(ir) - link!(ir, dyn_ir) + link_relocatable!(ir, relocations, dyn_ir, + dyn_meta.relocations) changed = true dyn_entry_fn end @@ -292,7 +295,7 @@ const __llvm_initialized = Ref(false) if job.config.toplevel && job.config.libraries # load the runtime outside of a timing block (because it recurses into the compiler) if !uses_julia_runtime(job) - runtime = load_runtime(job) + runtime, runtime_relocs = load_runtime(job) end @tracepoint "Library linking" begin @@ -301,7 +304,8 @@ const __llvm_initialized = Ref(false) # GPU run-time library if !uses_julia_runtime(job) - @tracepoint "runtime library" link!(ir, runtime; only_needed=true) + @tracepoint "runtime library" link_relocatable!( + ir, relocations, runtime, runtime_relocs; only_needed=true) end end end @@ -334,13 +338,16 @@ const __llvm_initialized = Ref(false) finish_linked_module!(job, ir) - # Materialize isbits and Bool boxes; bake addresses for other objects. - portable = relocate_gvs!(ir, gv_to_value) - portable || mark_session_dependent!(job) + # Resolve early so optimization sees concrete values. + resolve_early = resolve_relocations && relocation_lowering(job) === :bake + if resolve_early + prune_dead_relocations!(ir, relocations) + bake_relocations!(ir, relocations) + end if job.config.optimize @tracepoint "optimization" begin - optimize!(job, ir; job.config.opt_level) + optimize!(job, ir, relocations; job.config.opt_level) # deferred codegen has some special optimization requirements, # which also need to happen _after_ regular optimization. @@ -362,6 +369,12 @@ const __llvm_initialized = Ref(false) end end + # Runtime linking during optimization can add relocations. + if resolve_early && !isempty(relocations) + prune_dead_relocations!(ir, relocations) + bake_relocations!(ir, relocations) + end + if job.config.cleanup @tracepoint "clean-up" begin @dispose pb=NewPMPassBuilder() begin @@ -375,6 +388,9 @@ const __llvm_initialized = Ref(false) end end + # Do not expose dead sites to loaders. + resolve_early || prune_dead_relocations!(ir, relocations) + # optimization may have replaced functions, so look the entry point up again entry = functions(ir)[entry_fn] @@ -405,7 +421,20 @@ const __llvm_initialized = Ref(false) if job.config.toplevel && job.config.validate @tracepoint "validation" begin - check_ir(job, ir) + check_ir(job, ir, relocations) + end + end + + # Collect cglobal loads as relocation records now (after validation, which wants to see + # the original loads), so the `:llvm`-level relocation metadata is complete for loaders + # and for consumers that apply it themselves. Not under `:bake`, whose contract is a + # fully-resolved `:llvm` result; its loads are collected and resolved in + # `prepare_execution!`, which remains an idempotent safety net for the other strategies + # (and for direct `emit_asm` callers). + if job.config.toplevel && relocation_lowering(job) !== :bake + @tracepoint "cglobal relocations" begin + collect_cglobal_relocations!(job, ir, relocations) + prune_dead_relocations!(ir, relocations) end end @@ -413,21 +442,26 @@ const __llvm_initialized = Ref(false) @tracepoint "verification" verify(ir) end - return ir, (; entry, compiled, gv_to_value) + return ir, (; entry, compiled, relocations) end +# Compatibility for back-ends that resolve relocations during `emit_llvm`. +emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module, + format::LLVM.API.LLVMCodeGenFileType) = + emit_asm(job, ir, Relocations(), format) + @locked function emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module, - format::LLVM.API.LLVMCodeGenFileType) + relocs::Relocations, format::LLVM.API.LLVMCodeGenFileType) # NOTE: strip after validation to get better errors if job.config.strip @tracepoint "Debug info removal" strip_debuginfo!(ir) end @tracepoint "LLVM back-end" begin - @tracepoint "preparation" prepare_execution!(job, ir) + @tracepoint "preparation" prepare_execution!(job, ir, relocs) code = @tracepoint "machine-code generation" mcgen(job, ir, format) end - return code, () + return code, (;) end diff --git a/src/interface.jl b/src/interface.jl index a2fdcd04..fc2f0b71 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -302,6 +302,63 @@ can_vectorize(@nospecialize(job::CompilerJob)) = false # Should emit PTLS lookup that can be relocated dump_native(@nospecialize(job::CompilerJob)) = false +""" + relocation_lowering(job) -> Symbol + +Select how a back-end lowers relocations. There is one strategy per class of platform +capability, distinguished by what the platform lets you do to code *after* compiling it: + +- `:bake` (the default): the compiler resolves every record during `emit_llvm`, embedding + session-local addresses. Nothing is left for the loader, but the result cannot persist + across sessions. +- `:patch`: the platform has a symbol table and writable data after load. The compiler + emits each slot as a named global *definition* (interior records stay `extinit` + definitions); the loader looks up each record's `name`, adds its `offset`, and writes the + resolved word ([`resolved_relocations`](@ref)). Fits CUDA (`cuModuleGetGlobal` + + `cuMemcpyHtoD`) and equally an ORC consumer caching objects rather than IR. +- `:table`: the platform gives the loader no access to loaded code at all. The compiler + rewrites every record into an indexed load from a table of words supplied as run-time + data, obtained through [`relocation_table_pointer`](@ref); the loader materializes + [`resolved_relocation_table`](@ref) and hands over its base address at dispatch. Fits + Metal (the table travels in the kernel state). + +Both relocatable strategies produce session-portable results, enabling persistent caching +(`can_persist_results`); they require [`supports_relocatable_ir`](@ref). Resolving +relocations permanently roots the referenced Julia values in the process (mirroring Julia's +own codegen), so loaders need no GC bookkeeping of their own. + +When [`supports_relocatable_ir`](@ref) is false, Julia may still embed unrecorded addresses. +A back-end may lower the records it does receive with any strategy, but the result remains +session-local regardless. + +A loader may hold several compiled functions in one symbol namespace without renaming +anything: record names are namespaced per compiler job, and where two objects do define one +record — a relocation-carrying runtime-library function keeps its own namespace in every +kernel it is linked into — `:patch` defines it weakly, so the definitions coalesce and +patching the survivor serves both. + +The manifest is frozen once lowered: it then describes emitted code, so adding, dropping or +reordering a record errors. Take a [`copy`](@ref) to work on one afterwards, as +[`apply_relocations!`](@ref) does. +""" +relocation_lowering(@nospecialize(job::CompilerJob)) = :bake + +""" + relocation_table_pointer(job, builder, fun) -> LLVM.Value + +Emit, at `builder`'s current position inside `fun`, a pointer to the base of this job's +relocation table: an array of host words the loader fills from +[`resolved_relocation_table`](@ref). Required by back-ends selecting the `:table` +[`relocation_lowering`](@ref) strategy. + +The pointer must dominate the whole function, so it has to derive from something available +on entry — a kernel-state argument, typically. +""" +relocation_table_pointer(@nospecialize(job::CompilerJob), builder::IRBuilder, + fun::LLVM.Function) = + error("The `:table` relocation lowering requires " * + "$(nameof(typeof(job.config.target))) to implement `relocation_table_pointer`") + # the Julia module to look up target-specific runtime functions in (this includes both # target-specific functions from the GPU runtime library, like `malloc`, but also # replacements functions for operations like `Base.sin`) @@ -378,6 +435,32 @@ else cache_owner(job.config.target, job.config.params, job.config.always_inline) end +struct SessionResultEntry + config::CompilerConfig + value::Any +end + +const session_results_cache = IdDict{CodeInstance,Vector{SessionResultEntry}}() +const session_results_lock = ReentrantLock() + +function session_results(::Type{V}, ci::CodeInstance, config::CompilerConfig) where {V} + Base.@lock session_results_lock begin + entries = get!(session_results_cache, ci) do + SessionResultEntry[] + end + for entry in entries + entry.config === config && entry.value isa V && return entry.value::V + end + value = V() + push!(entries, SessionResultEntry(config, value)) + return value + end +end + +function can_persist_results(@nospecialize(job::CompilerJob)) + supports_relocatable_ir() && relocation_lowering(job) in (:patch, :table) +end + """ cached_results(::Type{V}, job::CompilerJob) -> Union{Nothing,V} @@ -409,21 +492,13 @@ post-compile lookup in the example is guaranteed to succeed. To attach results w generating code — e.g. from an inference-only precompilation workload — run `precompile(job)` first. -Results are keyed by the *full* compiler job: its method instance, world age, and entire -`CompilerConfig` (two jobs differing only in, say, the kernel `name` get distinct -structs). Storage differs per Julia version, transparently to the back-end: - -- On Julia 1.11+, the struct lives on the `CodeInstance`'s `analysis_results` chain in - Julia's integrated code cache, partitioned by [`cache_owner`](@ref) and keyed by - config within the per-CI [`JobResults`](@ref) container. Method redefinition - invalidates the CI — and with it the attached results. Artifacts stored during - precompilation are serialized into the package image along with the CI: populate only - session-portable values (bytes, strings) during precompile workloads, and keep - session-local handles (device modules, pipeline objects) in fields that remain empty - until first use at run time. +Results are keyed by the method instance, world age, and full `CompilerConfig`. On Julia +1.11+, results persist with the `CodeInstance` only when the back-end selects `:patch` or +`:table` and [`supports_relocatable_ir`](@ref) is true. Other results use a session-local +store. Julia 1.10 always uses the session-local store. -- On Julia 1.10, the struct lives in a session-local store keyed by the foreign - `CodeInstance` and config. Nothing persists across sessions. +Persistent result structs may contain session-local handles, but back-ends must leave those +fields empty during precompilation. Thread safety: concurrent calls for the same job return the same struct, but GPUCompiler does not serialize back-end *compilation*; take a back-end lock around @@ -434,7 +509,7 @@ function cached_results end @static if HAS_INTEGRATED_CACHE """ - JobResults + PersistentJobResults Per-CodeInstance container mapping a `CompilerConfig` to the back-end's results struct. @@ -449,28 +524,28 @@ from package images. # `CompilerConfig` is an abstract UnionAll here. Keeping it in a tuple stored inline in a # Vector boxes the config again on every iteration. A non-isbits entry object is stored by # reference, so both fields are boxed once and hot-path scans allocate nothing. -struct JobResultEntry +struct PersistentResultEntry config::CompilerConfig value::Any end -mutable struct JobResults - entries::Vector{JobResultEntry} - JobResults() = new(JobResultEntry[]) +mutable struct PersistentJobResults + entries::Vector{PersistentResultEntry} + PersistentJobResults() = new(PersistentResultEntry[]) end -const cached_results_lock = ReentrantLock() +const persistent_results_lock = ReentrantLock() # NOTE: like `cache_owner`, specialized for the launch hot path (bounded number of # instantiations: one per back-end and results type). -function job_results(::Type{V}, ci::CodeInstance, config::CompilerConfig) where {V} - jr = CompilerCaching.results(JobResults, ci) - Base.@lock cached_results_lock begin - for entry in jr.entries +function persistent_results(::Type{V}, ci::CodeInstance, config::CompilerConfig) where {V} + results = CompilerCaching.results(PersistentJobResults, ci) + Base.@lock persistent_results_lock begin + for entry in results.entries entry.config === config && entry.value isa V && return entry.value::V end v = V() - push!(jr.entries, JobResultEntry(config, v)) + push!(results.entries, PersistentResultEntry(config, v)) return v end end @@ -478,7 +553,7 @@ end function cache_view(@nospecialize(job::CompilerJob)) # `cache_owner` is deliberately stored as Any on CompilerConfig: preserve that box in the # CacheView instead of re-specializing and re-boxing the immutable token for the ccall. - CompilerCaching.CacheView{Any,JobResults}(cache_owner(job), job.world) + CompilerCaching.CacheView{Any,PersistentJobResults}(cache_owner(job), job.world) end # Fetch the CodeInstance backing `job` from the integrated cache. On 1.14+, inference @@ -503,54 +578,20 @@ end function cached_results(::Type{V}, job::CompilerJob) where {V} ci = job_code_instance(job) ci === nothing && return nothing - return job_results(V, ci, job.config) -end - -## session-dependent results -# -# Some compilation results embed session-specific data: `relocate_gvs!` bakes absolute -# pointers into the IR of toplevel jobs that reference `julia.constgv` globals (except -# for slots it can materialize as session-portable device constants), and any -# artifact a back-end derives from that IR (metallib, SPIR-V, ...) inherits them. Such -# results must not survive into a package image, while remaining available for -# within-session lookups during the precompilation process itself. Julia wipes its own -# session-dependent CodeInstance state during serialization (staticdata.c); we -# approximate that with an `atexit` hook, which the runtime invokes *before* -# `jl_write_compiler_output`: right before the image is written, the entries of jobs -# marked session-dependent are deleted from their `JobResults` container, so a later -# session simply recompiles them. - -const session_dependent_jobs = Vector{CompilerJob}() -const session_dependent_lock = ReentrantLock() - -function mark_session_dependent!(@nospecialize(job::CompilerJob)) - ccall(:jl_generating_output, Cint, ()) == 1 || return - Base.@lock session_dependent_lock begin - if isempty(session_dependent_jobs) - atexit(wipe_session_dependent_results) - end - push!(session_dependent_jobs, job) - end - return -end - -function wipe_session_dependent_results() - Base.@lock session_dependent_lock begin - for job in session_dependent_jobs - ci = job_code_instance(job) - ci === nothing && continue - jr = CompilerCaching.results(JobResults, ci) - Base.@lock cached_results_lock begin - filter!(entry -> entry.config !== job.config, jr.entries) - end - end - empty!(session_dependent_jobs) + if can_persist_results(job) + return persistent_results(V, ci, job.config) + else + return session_results(V, ci, job.config) end - return end end # HAS_INTEGRATED_CACHE +# Public relocation interface. +@public Relocation, RelocationSiteKind, SlotSite, InteriorSite, Relocations +@public relocation_lowering, relocation_table_pointer +@public apply_relocations!, resolved_relocations, resolved_relocation_table +@public supports_relocatable_ir @public GPUCompilerCacheToken, cache_owner, cached_results # the method table to use diff --git a/src/irgen.jl b/src/irgen.jl index 9129e20c..57f81ae6 100644 --- a/src/irgen.jl +++ b/src/irgen.jl @@ -2,6 +2,7 @@ function irgen(@nospecialize(job::CompilerJob)) mod, compiled, gv_to_value = @tracepoint "emission" compile_method_instance(job) + relocations = collect_julia_value_relocations!(job, mod, gv_to_value) if job.config.entry_abi === :specfunc entry_fn = compiled[job.source].specfunc else @@ -43,18 +44,22 @@ function irgen(@nospecialize(job::CompilerJob)) end end - # sanitize global values (Julia doesn't when using the external codegen policy) - for val in [collect(globals(mod)); collect(functions(mod))] + # sanitize defined globals (external declaration names are part of their ABI) + for val in collect(globals(mod)) + isdeclaration(val) && continue + new_name = safe_name(LLVM.name(val)) + if LLVM.name(val) != new_name + LLVM.name!(val, new_name) + end + end + + # sanitize defined functions (Julia doesn't when using the external codegen policy) + for val in collect(functions(mod)) isdeclaration(val) && continue old_name = LLVM.name(val) new_name = safe_name(old_name) if old_name != new_name LLVM.name!(val, new_name) - val = get(gv_to_value, old_name, nothing) - if val !== nothing - delete!(gv_to_value, old_name) - gv_to_value[new_name] = val - end end end @@ -145,7 +150,7 @@ function irgen(@nospecialize(job::CompilerJob)) lower_alloca!(job, mod) end - return mod, compiled, gv_to_value + return mod, compiled, relocations end @@ -539,6 +544,19 @@ end KERNEL_STATE # the kernel state argument end +# Compare an LLVM type converted from a Julia type against one read off a kernel signature, +# ignoring pointer address spaces. Post-optimization the two legitimately disagree there: +# `RemoveJuliaAddrspacesPass` strips Julia's tracked/derived address spaces during +# optimization, and a back-end may then move parameters into a target address space (Metal's +# device space, say), while `convert(LLVMType, T; allow_boxed=true)` keeps reporting +# `addrspace(10)`. Everything else about the shape still has to match. +function same_type_modulo_addrspaces(@nospecialize(a::LLVMType), @nospecialize(b::LLVMType)) + a == b && return true + (a isa LLVM.PointerType && b isa LLVM.PointerType) || return false + supports_typed_pointers(context()) || return true + return same_type_modulo_addrspaces(eltype(a), eltype(b)) +end + # Determine the calling convention of a the arguments of a Julia function, given the # LLVM function type as generated by the Julia code generator. Returns an vector with one # element for each Julia-level argument, containing a tuple with the following fields: @@ -586,7 +604,11 @@ function classify_arguments(@nospecialize(job::CompilerJob), codegen_ft::LLVM.Fu # - boxed values # XXX: use `deserves_retbox` instead? elseif llvm_source_typ isa LLVM.PointerType - @assert llvm_source_typ == codegen_typ + # a boxed argument is `addrspace(10)`-tracked when Julia emits it, but lands + # in whatever address space the back-end picked once optimization has run + @assert llvm_source_typ == codegen_typ || + (post_optimization && + same_type_modulo_addrspaces(llvm_source_typ, codegen_typ)) push!(args, (cc=MUT_REF, typ=source_typ, name=source_name, idx=codegen_i)) # - references to aggregates else diff --git a/src/jlgen.jl b/src/jlgen.jl index af94770f..0da33741 100644 --- a/src/jlgen.jl +++ b/src/jlgen.jl @@ -612,14 +612,10 @@ function compile_method_instance(@nospecialize(job::CompilerJob)) end end - # Maintain a map from global variables to their initialized Julia values. The - # objects pointed to are perma-rooted during codegen. We *don't* bake these - # addresses into the IR yet, so that we can cache it across sessions. + # Map global variables to their rooted Julia values without embedding addresses in IR. gv_to_value = Dict{String, Ptr{Cvoid}}() if gvs === nothing - # No reliable GV table on this Julia — best-effort discovery from the module. - # On these older versions Julia's own emit may have already baked in absolute - # pointer values; we recover them by reading existing initializers. + # Without a reliable GV table, recover addresses from existing initializers. for gv in globals(llvm_mod) if !haskey(metadata(gv), "julia.constgv") continue @@ -646,15 +642,8 @@ function compile_method_instance(@nospecialize(job::CompilerJob)) gv = GlobalVariable(gv_ref) gv_to_value[LLVM.name(gv)] = init end - # Strip the initializers so the IR we hand back is session-portable - # (on 1.12 `jl_emit_native_impl` bakes pointers via - # `literal_static_pointer_val`; on 1.13+ Julia nulls them itself); - # `relocate_gvs!` at the toplevel link step writes the session-current - # value in. Demote each GV to an external *declaration* rather than - # giving it a null initializer: an internal global initialized to null - # is fair game for optimization passes that run before relocation - # (e.g. GlobalOpt in the PTX back-end's `finish_module!`), which would - # fold loads of it to null and delete the global. + # Discard session addresses. Declarations survive optimization until relocation + # lowering; null definitions could be folded away first. for gv in globals(llvm_mod) haskey(gv_to_value, LLVM.name(gv)) || continue initializer!(gv, nothing) @@ -805,131 +794,6 @@ function compile_method_instance(@nospecialize(job::CompilerJob)) return llvm_mod, compiled, gv_to_value end -""" - relocate_gvs!(mod::LLVM.Module, gv_to_value::Dict{String, Ptr{Cvoid}}) - -Resolve globals that refer to Julia objects. `jl_true`/`jl_false`, which are -absent from `gv_to_value`, become canonical module-local boxes. Ordinary -non-ghost isbits objects are also materialized; other objects keep their host -address as an opaque identity token. - -Only GVs that are declarations (as produced by `compile_method_instance`, -which strips the initializers for session portability) or still have a null -initializer are touched: preserving existing initializers covers the -older-Julia path where Julia itself emits pointer values directly. - -Apart from the dedicated Bool globals, GVs present in `mod` but missing from -`gv_to_value` remain declarations, which back-ends will reject loudly -(undefined symbol) rather than silently folding to null. - -Returns `true` if the module is session-portable afterwards: no absolute host -address was written (neither a baked slot nor a materialized header carrying -a non-smalltag type pointer). -""" -function relocate_gvs!(mod::LLVM.Module, gv_to_value::Dict{String, Ptr{Cvoid}}) - portable = materialize_bool_singletons!(mod) - mod_gvs = globals(mod) - for (name, init) in gv_to_value - # Bools are resolved by name above. - name in ("jl_true", "jl_false") && continue - - haskey(mod_gvs, name) || continue - gv = mod_gvs[name] - cur = initializer(gv) - if !(cur === nothing || LLVM.isnull(cur)) - # pre-baked by Julia itself (pre-1.13): also session-absolute - portable = false - continue - end - - val = nothing - if init != C_NULL - obj = Base.unsafe_pointer_to_objref(init) - # Zero-sized objects remain identity tokens. - if isbitstype(typeof(obj)) && sizeof(obj) > 0 && !(obj isa Bool) - val, hdr = materialize_box!(mod, gv, obj, init) - # non-smalltag headers carry a host DataType pointer - portable &= hdr < UInt(64 << 4) # jl_max_tags << 4 - end - end - if val === nothing - val = const_inttoptr(ConstantInt(Int64(init)), global_value_type(gv)) - portable = false - end - initializer!(gv, val) - # re-internalize what compile_method_instance demoted to an external - # declaration; with the value in place, the optimizer can now fold it - linkage!(gv, LLVM.API.LLVMPrivateLinkage) - end - return portable -end - -# Bool JuliaVariables are absent from `gv_to_value`; define one device box per name. -function materialize_bool_singletons!(mod::LLVM.Module) - portable = true - mod_gvs = globals(mod) - for (name, obj) in ("jl_true" => true, "jl_false" => false) - haskey(mod_gvs, name) || continue - gv = mod_gvs[name] - cur = initializer(gv) - if !(cur === nothing || LLVM.isnull(cur)) - # Existing definitions may contain session-specific addresses. - portable = false - continue - end - - init = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), obj) - val, hdr = materialize_box!(mod, gv, obj, init) - initializer!(gv, val) - constant!(gv, true) - linkage!(gv, LLVM.API.LLVMPrivateLinkage) - - # Stay conservative if Bool stops using a smalltag. - portable &= hdr < UInt(64 << 4) # jl_max_tags << 4 - end - return portable -end - -# emit a device-resident constant replica of the box holding `obj`; returns -# the constant to store in the slot, and the (gcbits-masked) header word -function materialize_box!(mod::LLVM.Module, gv::GlobalVariable, @nospecialize(obj), - init::Ptr{Cvoid}) - @assert isbitstype(typeof(obj)) && sizeof(obj) > 0 - - W = sizeof(Int) - hdr, bytes = GC.@preserve obj begin - # the header word transparently yields the smalltag immediate for - # smalltag types and the host type pointer otherwise; drop the gcbits - hdr = unsafe_load(Ptr{UInt}(init - W)) & ~UInt(15) - bytes = [unsafe_load(Ptr{UInt8}(init), i) for i in 1:sizeof(obj)] - hdr, bytes - end - - T_word = LLVM.IntType(8W) - T_byte = LLVM.Int8Type() - fields = LLVM.Constant[ConstantInt(T_word, hdr), ConstantDataArray(T_byte, bytes)] - payload_idx = 1 - if Base.datatype_alignment(typeof(obj)) > W - # pad so the payload lands at a 16-byte offset (JL_HEAP_ALIGNMENT max) - pushfirst!(fields, ConstantDataArray(T_byte, zeros(UInt8, 16 - W))) - payload_idx = 2 - end - boxinit = ConstantStruct(fields) - boxty = value_type(boxinit) - - box = GlobalVariable(mod, boxty, safe_name(LLVM.name(gv)) * "_box") - initializer!(box, boxinit) - constant!(box, true) - linkage!(box, LLVM.API.LLVMPrivateLinkage) - alignment!(box, 16) - unnamed_addr!(box, true) - - idx(i) = ConstantInt(LLVM.Int32Type(), i) - payload = const_gep(boxty, box, LLVM.Constant[idx(0), idx(payload_idx)]) - slotty = global_value_type(gv) - val = value_type(payload) == slotty ? payload : const_addrspacecast(payload, slotty) - return val, hdr -end # partially revert JuliaLang/julia#49391 — see #527 @static if v"1.11.0-DEV.1603" <= VERSION < v"1.12.0-DEV.347" && # reverted on master diff --git a/src/mcgen.jl b/src/mcgen.jl index 3fb4a5c6..7f108617 100644 --- a/src/mcgen.jl +++ b/src/mcgen.jl @@ -1,79 +1,48 @@ # machine code generation -# final preparations for the module to be compiled to machine code -# these passes should not be run when e.g. compiling to write to disk. -function prepare_execution!(@nospecialize(job::CompilerJob), mod::LLVM.Module) - @dispose pb=NewPMPassBuilder() begin - register!(pb, ResolveCPUReferencesPass(job)) - - add!(pb, RecomputeGlobalsAAPass()) - add!(pb, GlobalOptPass()) - add!(pb, ResolveCPUReferencesPass(job)) - add!(pb, GlobalDCEPass()) - add!(pb, StripDeadPrototypesPass()) - - run!(pb, mod, llvm_machine(job.config.target)) - end - - return -end - -# some Julia code contains references to objects in the CPU run-time, -# without actually using the contents or functionality of those objects. -# -# prime example are type tags, which reference the address of the allocated type. -# since those references are ephemeral, we can't eagerly resolve and emit them in the IR, -# but at the same time the GPU can't resolve them at run-time. -# -# this pass performs that resolution at link time. -struct ResolveCPUReferences - job::CompilerJob -end -function (self::ResolveCPUReferences)(mod::LLVM.Module) - changed = false - - for f in functions(mod) - fn = LLVM.name(f) - if isdeclaration(f) && !LLVM.isintrinsic(f) && startswith(fn, "jl_") - # lazily resolve the address of the binding; some symbols only exist - # within the JIT (e.g. `jl_get_pgcstack_resolved`) and cannot be looked up, - # but such symbols are only ever called, not loaded from. - dereferenced = nothing - function resolve_binding() - if dereferenced === nothing - address = ccall(:jl_cglobal, Any, (Any, Any), fn, UInt) - dereferenced = LLVM.ConstantInt(unsafe_load(address)) - end - dereferenced - end - - function replace_bindings!(value) - changed = false - for use in uses(value) - val = user(use) - if isa(val, LLVM.ConstantExpr) - # recurse - changed |= replace_bindings!(val) - elseif isa(val, LLVM.LoadInst) - # resolve - replace_uses!(val, resolve_binding()) - erase!(val) - # FIXME: iterator invalidation? - changed = true - end +# Finalize the module for backend emission by collecting and lowering all live relocations. +function prepare_execution!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations=Relocations()) + # Clean up first so only live relocations get lowered. + function cleanup(; fold_instructions=false) + @dispose pb=NewPMPassBuilder() begin + if fold_instructions + add!(pb, NewPMFunctionPassManager()) do fpm + add!(fpm, instcombine_pass(job)) end - changed end - - changed |= replace_bindings!(f) + add!(pb, RecomputeGlobalsAAPass()) + add!(pb, GlobalOptPass()) + add!(pb, GlobalDCEPass()) + add!(pb, StripDeadPrototypesPass()) + run!(pb, mod, llvm_machine(job.config.target)) end end - - return changed + cleanup() + prune_dead_relocations!(mod, relocs) + + # For non-`:bake` strategies this already ran at the end of `emit_llvm` (so the + # `:llvm`-level metadata is complete); re-running is a no-op since rewritten loads + # target namespaced `gpu_jl_*` slots, which are not collection candidates. It remains + # load-bearing for `:bake` and for direct `emit_asm` callers that pass fresh + # `Relocations`. + collect_cglobal_relocations!(job, mod, relocs) + + # Lower, then freeze: from here the manifest describes emitted code, so no record may be + # added, dropped or reordered. Nothing can die on its own either — a record that survived + # the prune above is either baked into an initializer, anchored in `llvm.used`, or + # rewritten into a table load. Freezing after lowering, not before, leaves + # `bake_relocations!` free to consume the records as its final act. + lower_relocations!(job, mod, relocs) + freeze!(relocs) + + # Fold constants exposed by eager lowering, and drop globals the lowering left dead. + cleanup(; fold_instructions=true) + + has_unresolved_cglobal_loads(mod, relocs) && + error("Unresolved cglobal load after relocation lowering") + return end -ResolveCPUReferencesPass(job) = - NewPMModulePass("ResolveCPUReferences", ResolveCPUReferences(job)) - function mcgen(@nospecialize(job::CompilerJob), mod::LLVM.Module, format=LLVM.API.LLVMAssemblyFile) tm = llvm_machine(job.config.target) diff --git a/src/metal.jl b/src/metal.jl index 2e601083..bc73e5e2 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -187,6 +187,163 @@ function promote_bf16_intrinsics!(mod::LLVM.Module) return changed end +# Does `gv` look like a boxed-constant replica created by `materialize_box!`? Its `_box` name +# and layout distinguish it from other private constants: `{ i64 header, [payload bytes] }`, +# optionally with leading alignment padding. Relocatable boxes are external and mutable, so +# `isconstant` excludes them (they are delivered through the relocation table instead). +function is_boxed_constant(@nospecialize(gv::LLVM.GlobalVariable)) + isdeclaration(gv) && return false + endswith(LLVM.name(gv), "_box") || return false + linkage(gv) == LLVM.API.LLVMPrivateLinkage || return false + (isconstant(gv) && unnamed_addr(gv)) || return false + addrspace(value_type(gv)) == 0 || return false + T = global_value_type(gv) + T isa LLVM.StructType || return false + els = elements(T) + is_hdr(t) = t isa LLVM.IntegerType && width(t) == 8sizeof(UInt) + is_pad(t) = t isa LLVM.ArrayType && eltype(t) == LLVM.Int8Type() + return (length(els) >= 1 && is_hdr(els[1])) || + (length(els) >= 2 && is_pad(els[1]) && is_hdr(els[2])) +end + +# Are all (transitive constant) users of `gv` instructions, i.e. does the box escape only +# through function bodies? A box reachable from another global's *initializer* is the +# `jl_true`/`jl_false` slot→box indirection, which lives fine in the constant space +# (`add_global_address_spaces!`); demoting it would strand the slot's constant pointer. Only +# the isbits-union interior boxes — whose payload address flows through a body `phi`/`select` +# — need demotion, and those have instruction users exclusively. +function box_used_only_by_instructions(@nospecialize(gv::LLVM.GlobalVariable)) + ok = true + function walk(@nospecialize(v)) + for use in uses(v) + u = user(use) + if u isa LLVM.Instruction + # a body use: fine + elseif u isa LLVM.GlobalVariable + ok = false # initializer reference (slot→box) + elseif u isa LLVM.Constant + walk(u) + else + ok = false + end + end + end + walk(gv) + return ok +end + +# Demote materialized boxed-constant globals to per-function stack allocas. +# +# Motivation: an isbits `Union` return is lowered as `{ptr, i8}` whose payload pointer +# `phi`/`select`s the box global against the caller's `sret` alloca. AIR has no generic +# address space — AS 0 *is* thread memory — so once `add_global_address_spaces!` sinks the +# box into AS 2 (constant), the `addrspacecast` back to AS 0 at the use only works where the +# back-end can statically fold it away; across the union's call return or aggregate phi it +# cannot, and the load silently reads thread memory instead of the constant. Copying the box +# to a thread `alloca` is the sanctioned lowering (MSL rejects mixing `thread` and `constant` +# pointers outright); it is sound because box addresses carry no identity (isbits egal is by +# content) and Metal fully inlines device functions. +# +# Runs from `finish_ir!`: post-opt (the box only escapes its `materialize_box!` pointer slot +# once GlobalOpt folds it — still-slotted boxes like `jl_true`/`jl_false` are skipped and +# stay in constant space) and pre-`add_global_address_spaces!`. Relocatable boxes (external +# + `extinit`) are demoted by the `:table` lowering instead, which also fills their header. +function demote_boxed_constants!(mod::LLVM.Module) + changed = false + for gv in collect(globals(mod)) + (is_boxed_constant(gv) && box_used_only_by_instructions(gv)) || continue + if LLVM.version() < v"17" + # `replace_global_with_local!` needs LLVM.jl's `convert_users_to_instructions!` + # (LLVM 17+). Without demotion the kernel would compile but read thread memory + # instead of the boxed constant, so refuse loudly rather than miscompile. + error("Metal kernels embedding boxed union constants require Julia 1.12 or later") + end + boxty = global_value_type(gv) + init = initializer(gv) + # keep the box's alignment, but at least Julia's heap alignment (16 B), which is what + # `materialize_box!` gives these boxes and what the payload's `isbits` layout assumes + align = max(alignment(gv), 16) + slots = Dict{LLVM.Function, LLVM.Value}() + function slot(f::LLVM.Function) + get!(slots, f) do + @dispose builder=IRBuilder() begin + position!(builder, first(instructions(first(blocks(f))))) + ptr = alloca!(builder, boxty) + alignment!(ptr, align) + store!(builder, init, ptr) + ptr + end + end + end + replace_global_with_local!(gv, slot) + changed = true + end + return changed +end + + +## relocations as a kernel-state word table +# +# Metal gives a loader no access at all to loaded code: there is no post-load symbol patching +# (no ORC `absoluteSymbols`, no writable program-scope globals to `:patch`). So relocation words +# are delivered as ordinary *run-time data*: the loader resolves them in its session, writes +# `resolved_relocation_table` into a small buffer, and passes that buffer's device address in the +# kernel state at every dispatch. GPUCompiler's `:table` lowering rewrites each record into an +# indexed load off the base this hook returns. +# +# Nothing session-local ends up in the metallib, so it is byte-stable even for +# relocation-carrying kernels — which is what makes them persistable across sessions and +# content-keyable by `MTLBinaryArchive`. +# +# The contract with the back-end is one kernel-state field: +# +# reloc_table::Core.LLVMPtr{UInt64, AS.Device} +# +# always present (null for relocation-free kernels, which never read it) so that the state +# layout does not depend on what a kernel happens to reference. +const RELOCATION_TABLE_FIELD = :reloc_table + +function relocation_table_pointer(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + builder::IRBuilder, fun::LLVM.Function) + state = kernel_state_type(job) + field = state === Nothing ? nothing : + findfirst(isequal(RELOCATION_TABLE_FIELD), fieldnames(state)) + field === nothing && + error("""Metal delivers relocations through the kernel state, so its type must have a + `$(RELOCATION_TABLE_FIELD)::Core.LLVMPtr{UInt64, 1}` field; got $state.""") + + # The state arrives as the leading by-reference argument (`kernel_state_to_reference!`, + # then `add_parameter_address_spaces!`), so it is available on entry and dominates every + # use. Only a kernel has one, and Metal fully inlines device code, so only the kernel can + # be holding a relocation by now; insist on that rather than mistaking some other + # function's first pointer argument for the state. + (job.config.kernel && fun in kernels(LLVM.parent(fun))) || + error("""Metal delivers relocations through the kernel state, which only a kernel has. + Function `$(LLVM.name(fun))` is not one, so it cannot carry relocations; + compile it as a kernel, or select the `:bake` lowering.""") + state_ptr = parameters(fun)[1] + value_type(state_ptr) isa LLVM.PointerType || + error("Expected a kernel-state pointer argument, got $(value_type(state_ptr))") + + # Reach the field by its Julia byte offset rather than by a struct element index: Julia + # renders a struct as an LLVM array when all its fields share a type, and inserts explicit + # padding elements when they don't, so no element numbering matches the Julia one. + T_byte = LLVM.Int8Type() + T_table = convert(LLVMType, fieldtype(state, field)) + as = addrspace(value_type(state_ptr)) + typed = supports_typed_pointers(context()) + + base = typed ? bitcast!(builder, state_ptr, LLVM.PointerType(T_byte, as)) : state_ptr + field_ptr = inbounds_gep!(builder, T_byte, base, + [ConstantInt(LLVM.Int32Type(), fieldoffset(state, field))]) + typed && (field_ptr = bitcast!(builder, field_ptr, LLVM.PointerType(T_table, as))) + + table = load!(builder, T_table, field_ptr, "reloc_table") + alignment!(table, sizeof(UInt)) + return table +end + + function finish_linked_module!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module) # propagate `target.fastmath` as `@fastmath`-everywhere semantics, so the math-intrinsic # lowering in `finish_ir!` picks the relaxed `air.fast_*` functions. done here (post-link, @@ -447,6 +604,11 @@ function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L # add kernel metadata if job.config.kernel + # demote escaped boxed-constant globals to stack allocas before the address-space + # passes move constants into AS 2 (which would produce downgrade-incompatible IR for + # the ones whose address escapes an isbits union return; see above) + demote_boxed_constants!(mod) + entry = add_parameter_address_spaces!(job, mod, entry) entry = add_global_address_spaces!(job, mod, entry) @@ -548,12 +710,103 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L return end +# Julia names each generated LLVM function `julia__`, where the counter is +# drawn from a process-global codegen sequence and so differs from one session to the next. +# Left anywhere in the emitted bitcode it makes the AIR (and the metallib wrapping it) +# non-reproducible across sessions, defeating byte-stable caching and content-keyed binary +# archives. And a symbol name is not the only place it appears: inlining a Julia function +# leaves its name behind in the block labels the inliner synthesizes, in the (now orphaned) +# `DISubprogram` its debug locations still point at, and in the alias-scope strings Julia's +# codegen derived from it. +# +# So rewrite every occurrence, wherever it appears: map each distinct codegen name to a +# deterministic module-local form (its rank in a fixed traversal), then substitute that map +# into symbol names, value names, and metadata strings. A subprogram belonging to a function +# instead adopts that function's current name, so the entry — already renamed to a stable +# mangled symbol in `irgen.jl` — keeps a `linkageName` matching the symbol it describes. +# Runs on the final (post-`lower_air!`) module, just before the bitcode goes to the downgrader. +function normalize_julia_symbol_names!(mod::LLVM.Module) + # The counter is the trailing digit group, so match the name part lazily. The word + # boundaries keep the pattern from biting into a longer token: `myjulia_foo_1` is a user + # symbol that merely ends this way, and `julia_foo_12bar` is one that merely contains it. + # The name part must cover Julia's full method-name alphabet, not just `\w`: mutating + # functions carry `!` (`julia_record_exception!_18521`) and closures carry `#` + # (`julia_#kernel#123_456`), and a missed name leaks the per-session counter into the + # bitcode — exactly the byte-instability this pass exists to prevent. + codegen_name = r"(? "") * "_$(length(renames) + 1)" + end + normalize(str::AbstractString) = replace(str, codegen_name => deterministic) + rename!(value) = let str = LLVM.name(value) + occursin(codegen_name, str) && LLVM.name!(value, normalize(str)) + end + + # Metadata forms a graph (a debug location points at its subprogram, an alias scope at its + # domain), so walk it, replacing every string that carries a name. Only nodes reachable + # from a function or an instruction are visited, which is where inlining leaves its traces. + visited = Set{LLVM.API.LLVMMetadataRef}() + function normalize_metadata!(@nospecialize(md), replacement=nothing) + md isa LLVM.MDNode || return + md.ref in visited && return + push!(visited, md.ref) + for (i, op) in enumerate(operands(md)) + if op isa LLVM.MDString + str = convert(String, op) + occursin(codegen_name, str) || continue + new = replacement === nothing ? normalize(str) : replacement + new == str || LLVM.replace_operand(md, i, LLVM.MDString(new)) + elseif op isa LLVM.MDNode + normalize_metadata!(op) + end + end + end + + # the instruction metadata that can name a function: a debug location points at the + # subprogram it came from (orphaned once that function is inlined away), and Julia's alias + # scopes are labelled with the function they were derived for + md_kinds = (LLVM.MD_dbg, LLVM.MD_alias_scope, LLVM.MD_noalias, LLVM.MD_tbaa, + LLVM.MD_tbaa_struct, LLVM.MD_loop) + + # Symbols first, so that a subprogram can adopt its function's final name... + for f in functions(mod) + isdeclaration(f) || rename!(f) + end + # ...and every surviving function adopts before any instruction walk runs, since a walk + # can reach another function's subprogram through an inlined debug location and would + # otherwise give it a rank name instead of that function's symbol. + for f in functions(mod) + isdeclaration(f) && continue + sp = LLVM.subprogram(f) + sp === nothing || normalize_metadata!(sp, LLVM.name(f)) + end + for f in functions(mod) + isdeclaration(f) && continue + for bb in blocks(f) + rename!(bb) + for inst in instructions(bb) + rename!(inst) + md = metadata(inst) + for kind in md_kinds + haskey(md, kind) && normalize_metadata!(md[kind]) + end + end + end + end + return +end + @unlocked function mcgen(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module, format=LLVM.API.LLVMObjectFile) # lower LLVM constructs that the AIR back-end does not support; this takes the place # of instruction selection, as our LLVM does not have a Metal target machine. lower_air!(job, mod) + # scrub process-global Julia codegen counters from symbol / debug names, so the emitted + # bitcode is reproducible across sessions + normalize_julia_symbol_names!(mod) + if !isavailable(LLVMDowngrader_jll) error("Metal machine-code generation requires the LLVMDowngrader_jll package, which should be installed and loaded first.") end @@ -594,13 +847,22 @@ function add_parameter_address_spaces!(@nospecialize(job::CompilerJob), mod::LLV ft = function_type(f) # find the byref parameters - byref = BitVector(undef, length(parameters(ft))) + byref = falses(length(parameters(ft))) + # ... and among those, the boxed ones: an argument that survived `check_invocation` + # despite not being a bitstype has no fields, so it can only be used by identity (an + # interned `Symbol` being the typical case). Rather than a buffer, the host passes its + # address as a bare word, which this pass turns back into the pointer the body expects. + identity_word = falses(length(parameters(ft))) args = classify_arguments(job, ft; post_optimization=job.config.optimize) filter!(args) do arg arg.cc != GHOST end for arg in args - byref[arg.idx] = (arg.cc == BITS_REF || arg.cc == KERNEL_STATE) + param = parameters(ft)[arg.idx] + identity_word[arg.idx] = arg.cc == MUT_REF && param isa LLVM.PointerType && + addrspace(param) == 0 + byref[arg.idx] = arg.cc == BITS_REF || arg.cc == KERNEL_STATE || + identity_word[arg.idx] end function remapType(src) @@ -648,7 +910,16 @@ function add_parameter_address_spaces!(@nospecialize(job::CompilerJob), mod::LLV # perform argument conversions for (i, param) in enumerate(parameters(ft)) - if byref[i] + if identity_word[i] + # recover the boxed argument's address from the word the host passed + T_word = convert(LLVMType, UInt) + slot = parameters(new_f)[i] + if supports_typed_pointers(context()) + slot = bitcast!(builder, slot, + LLVM.PointerType(T_word, addrspace(value_type(slot)))) + end + push!(new_args, inttoptr!(builder, load!(builder, T_word, slot), param)) + elseif byref[i] # load the argument in a stack slot llvm_typ = convert(LLVMType, args[i].typ) val = load!(builder, llvm_typ, parameters(new_f)[i]) @@ -1334,7 +1605,11 @@ function add_argument_metadata!(@nospecialize(job::CompilerJob), mod::LLVM.Modul push!(md, MDString("air.address_space")) push!(md, Metadata(ConstantInt(Int32(addrspace(parameters(entry_ft)[arg.idx]))))) - arg_type = if arg.typ <: Core.LLVMPtr + # A fieldless boxed argument (e.g. an interned `Symbol`) is passed as its bare address + # word, so describe that word: its Julia type has no size to report. + arg_type = if arg.cc == MUT_REF + UInt + elseif arg.typ <: Core.LLVMPtr arg.typ.parameters[1] else arg.typ diff --git a/src/optim.jl b/src/optim.jl index b0b339a3..932f316b 100644 --- a/src/optim.jl +++ b/src/optim.jl @@ -49,7 +49,8 @@ function aggressiveinstcombine_pass(@nospecialize(job::CompilerJob)) end end -function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level=2) +function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations; opt_level=2) tm = llvm_machine(job.config.target) tti = llvm_targetinfo(job.config.target) @@ -59,7 +60,7 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level= register!(pb, GPULowerCPUFeaturesPass(job)) register!(pb, GPULowerPTLSPass(job)) register!(pb, GPULowerGCFramePass(job)) - register!(pb, GPULinkRuntimePass(job)) + register!(pb, GPULinkRuntimePass(job, relocs)) register!(pb, GPULinkLibrariesPass(job)) register!(pb, GPUFinishRuntimeIntrinsicsPass(job)) register!(pb, AddKernelStatePass(job)) @@ -325,7 +326,8 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op add!(fpm, GPULowerGCFramePass(job)) end if job.config.libraries - add!(mpm, GPULinkRuntimePass(job)) + # Use the registered pass because it owns `relocs`; the others only capture `job`. + add!(mpm, "GPULinkRuntime") add!(mpm, GPULinkLibrariesPass(job)) add!(mpm, GPUFinishRuntimeIntrinsicsPass(job)) end @@ -475,6 +477,7 @@ GPULowerCPUFeaturesPass(job) = NewPMModulePass("GPULowerCPUFeatures", CPUFeature struct LinkRuntime job::CompilerJob + relocs::Relocations end function (self::LinkRuntime)(mod::LLVM.Module) self.job.config.libraries || return false @@ -483,15 +486,16 @@ function (self::LinkRuntime)(mod::LLVM.Module) # GC lowering can introduce new calls to GPU runtime functions after the runtime # was linked initially. Link again now so those calls resolve to definitions before # later intrinsic-lowering passes inspect or rewrite the runtime call graph. - runtime = load_runtime(self.job) + runtime, runtime_relocs = load_runtime(self.job) # `RemoveNIPass` stripped non-integral address spaces from `mod`'s datalayout, but the # cached runtime kept them; align it (as with target libraries) to avoid a warning. triple!(runtime, triple(mod)) datalayout!(runtime, datalayout(mod)) - link!(mod, runtime; only_needed=true) + link_relocatable!(mod, self.relocs, runtime, runtime_relocs; only_needed=true) return true end -GPULinkRuntimePass(job) = NewPMModulePass("GPULinkRuntime", LinkRuntime(job)) +GPULinkRuntimePass(job, relocs::Relocations) = + NewPMModulePass("GPULinkRuntime", LinkRuntime(job, relocs)) struct LinkLibraries job::CompilerJob diff --git a/src/relocation.jl b/src/relocation.jl new file mode 100644 index 00000000..6436688d --- /dev/null +++ b/src/relocation.jl @@ -0,0 +1,890 @@ +# Relocations name words in a module that hold a host address: a reference to a Julia value, +# or a word read from a named C global. Each is recorded as a typed [`Relocation`](@ref) whose +# `kind` says what the site *is*; no lowering infers that from the IR shape. +# +# produce ──▶ merge (on link) ──▶ prune (after DCE) ──▶ lower +# +# Site names are fixed at creation and namespaced by their producer, so IR and metadata can +# be linked without renaming. Unrelated jobs use distinct names; runtime functions linked +# into several outputs reuse names whose definitions and targets are identical. + + +## targets + +""" + JuliaValueRef(value) + +A Julia value with a stable address (a heap object, symbol, or singleton), used as the +serializable identity of a relocation target. Resolve it in the active session with +[`resolve_relocation_target`](@ref), which permanently roots the value in the process so +the resolved address stays valid for as long as the session lives. +""" +struct JuliaValueRef + value::Any + + function JuliaValueRef(value) + value_type = typeof(value) + isbitstype(value_type) && sizeof(value_type) > 0 && + error("JuliaValueRef requires an object with a stable address") + new(value) + end +end + +""" + CGlobalRef(symbol, library=nothing) + +A named C data global. With `library === nothing`, resolution uses `jl_cglobal`'s +process-wide lookup. Otherwise it looks up `symbol` in `library`. Resolution returns the +word stored in the global. +""" +struct CGlobalRef + symbol::Symbol + library::Union{Nothing,String} +end +CGlobalRef(symbol::Symbol) = CGlobalRef(symbol, nothing) + +""" + RelocationTarget + +A serializable target for a relocated word: either a [`JuliaValueRef`](@ref) or a +[`CGlobalRef`](@ref). +""" +const RelocationTarget = Union{JuliaValueRef,CGlobalRef} + +same_relocation_target(a::JuliaValueRef, b::JuliaValueRef) = a.value === b.value +same_relocation_target(a::CGlobalRef, b::CGlobalRef) = + a.symbol === b.symbol && a.library == b.library +same_relocation_target(::RelocationTarget, ::RelocationTarget) = false + +# Permanently root a value in the current process and return the canonical rooted +# instance, exactly as Julia's own codegen does for values referenced from native code +# (`jl_ensure_rooted` on 1.10/1.11, `aot_optimize_roots` on 1.12+). Values compiled in +# this session are already rooted this way, making this a cheap lookup; it matters for +# metadata deserialized from a cache, whose values codegen never saw. Rooting by egal +# identity also folds such duplicates onto the instance native code already uses. +function root_relocation_target(target::JuliaValueRef) + @static if VERSION >= v"1.11-" + ccall(:jl_as_global_root, Any, (Any, Cint), target.value, 1) + else + ccall(:jl_as_global_root, Any, (Any,), target.value) + end +end + +value_pointer(@nospecialize(value)) = UInt(ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), value)) + +""" + resolve_relocation_target(target) -> UInt + +Resolve a relocation target to its word in the current Julia process. Julia values are +permanently rooted (and canonicalized by egal identity) as part of resolution, so the +returned address cannot dangle. +""" +function resolve_relocation_target(target::JuliaValueRef) + value_pointer(root_relocation_target(target)) +end +function resolve_relocation_target(target::CGlobalRef) + if target.library === nothing + # `jl_cglobal` accepts the symbol directly and does the process-wide `jl_dlfind`. + address = ccall(:jl_cglobal, Any, (Any, Any), target.symbol, UInt) + return unsafe_load(address) + end + handle = Libdl.dlopen(target.library) + address = Libdl.dlsym(handle, target.symbol) + return unsafe_load(Ptr{UInt}(address)) +end + + +## the table + +""" + RelocationSiteKind + +What a relocation record points at, and hence how it is lowered: + +- `SlotSite`: a word-sized global the code loads through (GOT-style). Produced for + references to Julia values and for `cglobal` words. +- `InteriorSite`: a word inside a definition's initializer, namely the header of a + materialized box (see `materialize_box!`). +""" +@enum RelocationSiteKind SlotSite InteriorSite + +""" + Relocation(kind, name, offset, target) + +One word to relocate: a global `name` (unique to the job that produced it), a byte `offset` +within that global (always zero for a [`SlotSite`](@ref RelocationSiteKind)), the site +`kind`, and the [`RelocationTarget`](@ref) whose address belongs there. +""" +struct Relocation + kind::RelocationSiteKind + name::String + offset::Int + target::RelocationTarget + + function Relocation(kind::RelocationSiteKind, name::String, offset::Int, + target::RelocationTarget) + offset >= 0 || throw(ArgumentError("relocation offset must be nonnegative")) + kind === SlotSite && offset != 0 && + throw(ArgumentError("a relocation slot must have offset zero")) + new(kind, name, offset, target) + end +end + +# Records are kept sorted by this key, which is also their identity: at most one record per +# word. Ordering makes the record vector a deterministic manifest, which is what lets the +# `:table` lowering index the words by rank reproducibly. +relocation_key(rec::Relocation) = (rec.name, rec.offset) + +""" + Relocations(records) + +Relocation metadata accompanying a module: [`Relocation`](@ref) records sorted by +`(name, offset)`. See [`resolved_relocations`](@ref) and +[`resolved_relocation_table`](@ref) for handing them to a loader. + +Lowering is the end of the manifest's mutable life (`produce → merge → prune → lower → +freeze`): from then on it describes emitted code, which cannot be renegotiated, so the +mutators refuse to touch it. Work on a [`copy`](@ref) if you need a mutable one. +""" +struct Relocations + records::Vector{Relocation} + # The `:table` lowering's word order, materialized by it so that the delivered words + # cannot be desynced from the indices it baked into the code (see + # `emit_table_relocations!`). Empty for every other strategy. + table::Vector{RelocationTarget} + frozen::Base.RefValue{Bool} +end + +Relocations(records::Vector{Relocation}) = + Relocations(records, RelocationTarget[], Ref(false)) +Relocations() = Relocations(Relocation[]) + +# Resolving into IR consumes the records; loaders (and anything else working after lowering) +# copy cached metadata first, which is also how they get a mutable manifest back. +Base.copy(relocs::Relocations) = + Relocations(copy(relocs.records), copy(relocs.table), Ref(false)) +Base.isempty(relocs::Relocations) = isempty(relocs.records) +Base.length(relocs::Relocations) = length(relocs.records) + +# Lowering has committed the manifest to emitted code: adding, removing or reordering a +# record now silently desyncs it from that code — a `:table` index shifts onto the wrong +# word, a `:patch` definition is left holding a zero. Refuse instead. +freeze!(relocs::Relocations) = (relocs.frozen[] = true; relocs) + +function check_mutable(relocs::Relocations, what::String) + relocs.frozen[] && + error("""Cannot $what a relocation manifest that has already been lowered: its + records describe emitted code. Work on a `copy` instead.""") + return +end + +# Binary-search `records` for `key`: the index it occupies, or the index it would be +# inserted at, plus whether it is present. +function relocation_index(records::Vector{Relocation}, key::Tuple{String,Int}) + lo, hi = 1, length(records) + while lo <= hi + mid = (lo + hi) >>> 1 + found = relocation_key(records[mid]) + if found < key + lo = mid + 1 + elseif found > key + hi = mid - 1 + else + return mid, true + end + end + return lo, false +end + +# Record `rec`, keeping `records` sorted. A record for the same word must agree on +# everything but is otherwise accepted, so that linking two modules that both reference a +# value merges their metadata. +function add_relocation!(relocs::Relocations, rec::Relocation) + check_mutable(relocs, "add to") + records = relocs.records + idx, present = relocation_index(records, relocation_key(rec)) + if present + existing = records[idx] + same_relocation_target(existing.target, rec.target) || + error("Relocation '$(rec.name)+$(rec.offset)' refers to conflicting values") + existing.kind === rec.kind || + error("Relocation '$(rec.name)+$(rec.offset)' is recorded as both " * + "$(existing.kind) and $(rec.kind)") + return existing + end + insert!(records, idx, rec) + return rec +end + +add_relocation!(relocs::Relocations, kind::RelocationSiteKind, name::String, offset::Int, + target::RelocationTarget) = + add_relocation!(relocs, Relocation(kind, name, offset, target)) + +# The record for `(name, offset)`, or `nothing`. +function find_relocation(relocs::Relocations, name::String, offset::Int=0) + idx, present = relocation_index(relocs.records, (name, offset)) + return present ? relocs.records[idx] : nothing +end + +""" + resolved_relocations(relocs) -> Vector{Pair{Relocation,UInt}} + +Resolve relocation metadata for a `:patch` loader, returning each record with its resolved +word. Resolution permanently roots referenced Julia values in the process, so the addresses +stay valid for the lifetime of the session. +""" +function resolved_relocations(relocs::Relocations) + return Pair{Relocation,UInt}[rec => resolve_relocation_target(rec.target) + for rec in relocs.records] +end + +""" + resolved_relocation_table(relocs) -> Vector{UInt} + +Resolve relocation metadata for a `:table` loader, returning the words in the order the +`:table` lowering indexed them by. Resolution permanently roots referenced Julia values in +the process, so the addresses stay valid for the lifetime of the session. +""" +function resolved_relocation_table(relocs::Relocations) + isempty(relocs.table) && !isempty(relocs.records) && + error("""This manifest has $(length(relocs.records)) relocation record(s) but no + lowered table, so the code that reads it was never rewritten. Hand the + manifest to `emit_asm` (the 4-argument form) rather than emitting the + module with an empty one.""") + return UInt[resolve_relocation_target(target) for target in relocs.table] +end + +relocation_word_type() = LLVM.IntType(8sizeof(UInt)) + +function check_slot_size(mod::LLVM.Module, gv::GlobalVariable, name::String) + size = abi_size(datalayout(mod), global_value_type(gv)) + size == sizeof(UInt) || + error("Relocation slot '$name' has size $size, expected $(sizeof(UInt))") + return +end + +function slot_initializer(gv::GlobalVariable, value::UInt) + T = global_value_type(gv) + if T isa LLVM.PointerType + return const_inttoptr(ConstantInt(UInt64(value)), T) + elseif T isa LLVM.IntegerType && width(T) == 8sizeof(UInt) + return ConstantInt(T, value) + end + error("Relocation slot '$(LLVM.name(gv))' has unsupported LLVM type $T") +end + +# Validate `gv` against what `rec` says it is. The record is authoritative: a mismatch means +# the metadata and the IR have drifted apart, which every lowering would otherwise turn into +# a silently wrong word. +function check_relocation(mod::LLVM.Module, rec::Relocation, gv::GlobalVariable) + if rec.kind === SlotSite + check_slot_size(mod, gv, rec.name) + else + isdeclaration(gv) && + error("Interior relocation '$(rec.name)' is a declaration") + init = initializer(gv) + init === nothing && error("Relocation global '$(rec.name)' has no initializer") + T = value_type(init) + T isa LLVM.StructType || + error("Relocation global '$(rec.name)' has non-struct initializer $T") + size = abi_size(datalayout(mod), T) + rec.offset + sizeof(UInt) <= size || + error("Relocation '$(rec.name)+$(rec.offset)' is outside its $size-byte global") + end + return +end + +function foreach_relocation(f, mod::LLVM.Module, relocs::Relocations) + mod_gvs = globals(mod) + for rec in relocs.records + haskey(mod_gvs, rec.name) || error("Missing relocation global '$(rec.name)'") + gv = mod_gvs[rec.name] + check_relocation(mod, rec, gv) + f(rec, gv) + end + return +end + + +## producers + +# Julia names value globals `_`, where `` comes from a process-global +# codegen sequence and so differs from one session to the next. Drop it from relocation slot and +# box names: the target's `objectid` (appended alongside) is the stable per-target identity that +# disambiguates them, so any bitcode keyed on these names stays reproducible across sessions. +# `objectid` is content-stable for the interned symbols and `isbits`/`DataType` values that +# appear as relocation targets. +strip_codegen_counter(name::AbstractString) = replace(name, r"_[0-9]+$" => "") + +# A namespace for this job's relocation site names, so that no two kernels — nor a kernel and +# the runtime library linked into it — can ever define the same site symbol. That makes the +# names globally unique by construction, which is what lets `:patch` loaders share one symbol +# namespace (e.g. an ORC `JITDylib` holding several compiled functions) without renaming. +# +# The job's entry name is the natural discriminator, but it is only fixed later in `irgen`, and +# for an unnamed non-kernel job it would carry Julia's per-session codegen counter. Derive the +# same name deterministically instead: the configured name if there is one, otherwise the +# mangled signature (which is literally the entry name for kernels). +relocation_namespace(@nospecialize(job::CompilerJob)) = + job.config.name !== nothing ? safe_name(job.config.name) : + mangle_sig(job.source.specTypes) + +# Site names are used as symbols by loaders, so they must survive every back-end's assembler +# (`ptxas` in particular rejects anything outside `[A-Za-z0-9_$]`); `safe_name` guarantees that +# and `_` is the only separator available. +namespaced_name(namespace::String, base::AbstractString) = namespace * "_" * base + +function collect_julia_value_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + gv_to_value::Dict{String, Ptr{Cvoid}}) + relocs = Relocations() + namespace = relocation_namespace(job) + mod_gvs = globals(mod) + for (name, init) in gv_to_value + haskey(mod_gvs, name) || continue + gv = mod_gvs[name] + cur = initializer(gv) + if !(cur === nothing || LLVM.isnull(cur)) + @assert !supports_relocatable_ir() + continue + end + + # jl_get_llvm_gvs and jl_get_llvm_gv_inits report an initializer for every + # mapped global, so a null here means those maps are out of sync. + init == C_NULL && error("Missing Julia object for global '$name'") + obj = Base.unsafe_pointer_to_objref(init) + if isbitstype(typeof(obj)) && sizeof(typeof(obj)) > 0 && !(obj isa Bool) + val = materialize_box!(mod, relocs, namespace, gv, obj, init) + initializer!(gv, val) + linkage!(gv, LLVM.API.LLVMPrivateLinkage) + else + check_slot_size(mod, gv, name) + slot_name = namespaced_name(namespace, + strip_codegen_counter(safe_name(name)) * "_" * + string(objectid(obj); base=16)) + # Codegen can emit several slots for one value in a module (observed on 1.11, + # whose backported GV API does not deduplicate them), and their content-derived + # names collide by construction. Alias later slots onto the first: an equal name + # means an equal referenced value, and `add_relocation!` below degenerates into + # its agreeing-duplicate no-op (or errors on the astronomically unlikely + # `objectid` collision between distinct values). + existing = haskey(mod_gvs, slot_name) ? mod_gvs[slot_name] : nothing + if existing !== nothing && existing !== gv + @assert value_type(existing) == value_type(gv) + replace_uses!(gv, existing) + erase!(gv) + else + LLVM.name!(gv, slot_name) + LLVM.name(gv) == slot_name || + error("Relocation slot name '$slot_name' is already in use") + end + add_relocation!(relocs, SlotSite, slot_name, 0, JuliaValueRef(obj)) + end + end + + # Bool JuliaVariables are absent from `gv_to_value`; define one device box per module. + for (name, obj) in ("jl_true" => true, "jl_false" => false) + haskey(mod_gvs, name) || continue + gv = mod_gvs[name] + cur = initializer(gv) + if !(cur === nothing || LLVM.isnull(cur)) + @assert !supports_relocatable_ir() + continue + end + + init = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), obj) + val = materialize_box!(mod, relocs, namespace, gv, obj, init) + initializer!(gv, val) + constant!(gv, true) + linkage!(gv, LLVM.API.LLVMPrivateLinkage) + end + return relocs +end + +# Emit a device-resident constant replica of the box holding `obj` and return +# the constant to store in its slot. Any relocatable header is recorded in `relocs`. +function materialize_box!(mod::LLVM.Module, relocs::Relocations, namespace::String, + gv::GlobalVariable, @nospecialize(obj), init::Ptr{Cvoid}) + obj_type = typeof(obj) + @assert isbitstype(obj_type) + obj_size = sizeof(obj_type) + @assert obj_size > 0 + + W = sizeof(Int) + hdr, bytes = GC.@preserve obj begin + # the header word transparently yields the smalltag immediate for + # smalltag types and the host type pointer otherwise; drop the gcbits + hdr = unsafe_load(Ptr{UInt}(init - W)) & ~UInt(15) + bytes = [unsafe_load(Ptr{UInt8}(init), i) for i in 1:obj_size] + hdr, bytes + end + + T_word = LLVM.IntType(8W) + T_byte = LLVM.Int8Type() + patch_header = hdr >= UInt(64 << 4) # jl_max_tags << 4 + fields = LLVM.Constant[ConstantInt(T_word, patch_header ? 0 : hdr), + ConstantDataArray(T_byte, bytes)] + header_idx = 0 + payload_idx = 1 + if Base.datatype_alignment(typeof(obj)) > W + # pad so the payload lands at a 16-byte offset (JL_HEAP_ALIGNMENT max) + pushfirst!(fields, ConstantDataArray(T_byte, zeros(UInt8, 16 - W))) + header_idx = 1 + payload_idx = 2 + end + boxinit = ConstantStruct(fields) + boxty = value_type(boxinit) + + # Only a relocatable box needs a namespaced name: its header is a site loaders address by + # name. A fully-materialized box is a private constant, so LLVM uniques it on its own. + box_name = if patch_header + namespaced_name(namespace, + strip_codegen_counter(safe_name(LLVM.name(gv))) * "_" * + string(objectid(obj); base=16) * "_box") + else + safe_name(LLVM.name(gv)) * "_box" + end + box = GlobalVariable(mod, boxty, box_name) + LLVM.name(box) == box_name || error("Interior relocation global '$box_name' is already in use") + initializer!(box, boxinit) + alignment!(box, 16) + if patch_header + constant!(box, false) + linkage!(box, LLVM.API.LLVMExternalLinkage) + extinit!(box, true) + offset = Int(offsetof(datalayout(mod), boxty, header_idx)) + add_relocation!(relocs, InteriorSite, box_name, offset, JuliaValueRef(typeof(obj))) + else + constant!(box, true) + linkage!(box, LLVM.API.LLVMPrivateLinkage) + unnamed_addr!(box, true) + end + + idx(i) = ConstantInt(LLVM.Int32Type(), i) + payload = const_gep(boxty, box, LLVM.Constant[idx(0), idx(payload_idx)]) + slotty = global_value_type(gv) + val = value_type(payload) == slotty ? payload : const_addrspacecast(payload, slotty) + return val +end + +# Rewrite every load of `value` into the word `produce_word(builder)` emits at the load's +# position, restoring a pointer with `inttoptr` where the original load produced one. +# Constant expressions (typed-pointer bitcasts) are recursed through, so both `i64` and +# pointer-typed words are handled. Shared by every producer and lowering that replaces a +# word-sized global with a run-time value. +function rewrite_word_loads!(produce_word, @nospecialize(value), what::String) + changed = false + for use in collect(uses(value)) + val = user(use) + if isa(val, LLVM.ConstantExpr) + changed |= rewrite_word_loads!(produce_word, val, what) + elseif isa(val, LLVM.LoadInst) + T = value_type(val) + (T isa LLVM.PointerType || + (T isa LLVM.IntegerType && width(T) == 8sizeof(UInt))) || + error("Unsupported $what load of LLVM type $T") + @dispose builder=IRBuilder() begin + position!(builder, val) + replacement = produce_word(builder) + T isa LLVM.PointerType && + (replacement = inttoptr!(builder, replacement, T)) + replace_uses!(val, replacement) + end + erase!(val) + changed = true + end + end + return changed +end + +# Some Julia code loads words from libjulia C globals, for example type tags. Record those +# loads as dedicated zero-offset relocations immediately before object emission. +function is_cglobal_candidate(value, relocs::Relocations) + name = LLVM.name(value) + value isa LLVM.GlobalVariable && + find_relocation(relocs, name) !== nothing && return false + isdeclaration(value) || return false + value isa LLVM.Function && LLVM.isintrinsic(value) && return false + return startswith(name, "jl_") +end + +function collect_cglobal_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + changed = false + namespace = relocation_namespace(job) + + for f in [collect(functions(mod)); collect(globals(mod))] + is_cglobal_candidate(f, relocs) || continue + fn = LLVM.name(f) + slot = nothing + function cglobal_slot() + if slot === nothing + name = namespaced_name(namespace, "gpu_" * fn) + slot = GlobalVariable(mod, relocation_word_type(), name) + LLVM.name(slot) == name || + error("cglobal slot name '$name' is already in use") + add_relocation!(relocs, SlotSite, name, 0, CGlobalRef(Symbol(fn))) + end + slot + end + + changed |= rewrite_word_loads!(f, "cglobal '$fn'") do builder + load!(builder, relocation_word_type(), cglobal_slot()) + end + end + + return changed +end + +function has_unresolved_cglobal_loads(mod::LLVM.Module, relocs::Relocations) + function has_load(value) + for use in uses(value) + val = user(use) + val isa LLVM.LoadInst && return true + val isa LLVM.ConstantExpr && has_load(val) && return true + end + return false + end + + for value in [collect(functions(mod)); collect(globals(mod))] + is_cglobal_candidate(value, relocs) || continue + has_load(value) && return true + end + return false +end + + +## bookkeeping + +# Merge `src_mod` into `dest_mod` and carry its relocation metadata across. A site name always +# denotes the same word and target; duplicate declarations or definitions may therefore be +# coalesced by LLVM before their agreeing records are merged. +function link_relocatable!(dest_mod::LLVM.Module, dest_relocs::Relocations, + src_mod::LLVM.Module, src_relocs::Relocations; + only_needed=false) + link!(dest_mod, src_mod; only_needed) + for rec in src_relocs.records + # A site absent from the linked module was dead (DCE'd or not imported under + # `only_needed`); its relocation dies with it. + haskey(globals(dest_mod), rec.name) || continue + add_relocation!(dest_relocs, rec) + end + return +end + +function prune_dead_relocations!(mod::LLVM.Module, relocs::Relocations) + check_mutable(relocs, "prune") + mod_gvs = globals(mod) + dead_names = Set{String}() + for rec in relocs.records + gv = haskey(mod_gvs, rec.name) ? mod_gvs[rec.name] : nothing + if gv === nothing || (!isdeclaration(gv) && isempty(uses(gv))) + push!(dead_names, rec.name) + end + end + filter!(rec -> !(rec.name in dead_names), relocs.records) + for name in dead_names + gv = haskey(mod_gvs, name) ? mod_gvs[name] : nothing + gv === nothing || isdeclaration(gv) || erase!(gv) + end + return +end + + +## lowering + +# Lower live relocations before object emission, dispatching on the back-end's +# `relocation_lowering` strategy. Internal: back-ends select a strategy through the trait +# rather than overriding this. +function lower_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + strategy = relocation_lowering(job) + if strategy === :bake + bake_relocations!(mod, relocs) + elseif strategy === :patch + emit_patchable_relocations!(mod, relocs) + elseif strategy === :table + emit_table_relocations!(job, mod, relocs) + else + error("Unknown relocation lowering strategy :$strategy") + end + return +end + +# Overwrite the word at `offset` in `gv`'s struct initializer with `word`. +function patch_initializer_word!(mod::LLVM.Module, gv::GlobalVariable, offset::Int, + word::UInt) + init = initializer(gv) + T = value_type(init)::LLVM.StructType + idx = Int(element_at(datalayout(mod), T, offset)) + 1 + # An all-zero box (e.g. a patchable header over a zero payload) is folded to a + # `zeroinitializer`, a `ConstantAggregateZero` that reports no operands; rebuild + # the explicit per-field constants from the struct's element types. + fields = if init isa LLVM.ConstantAggregateZero + LLVM.Constant[null(elty) for elty in elements(T)] + else + LLVM.Constant[operands(init)...] + end + fields[idx] = ConstantInt(value_type(fields[idx]), word) + initializer!(gv, ConstantStruct(T, fields)) + return +end + +""" + bake_relocations!(mod, relocs) + +Resolve every record in the current Julia process and write the resulting words into the IR, +leaving `relocs` empty. The module then embeds session-local addresses and must not be +persisted across sessions. Drop dead records first with [`prune_dead_relocations!`](@ref). +""" +function bake_relocations!(mod::LLVM.Module, relocs::Relocations) + check_mutable(relocs, "resolve into IR") + foreach_relocation(mod, relocs) do rec, gv + word = resolve_relocation_target(rec.target) + if rec.kind === SlotSite + initializer!(gv, slot_initializer(gv, word)) + linkage!(gv, LLVM.API.LLVMPrivateLinkage) + constant!(gv, true) + else + patch_initializer_word!(mod, gv, rec.offset, word) + linkage!(gv, LLVM.API.LLVMPrivateLinkage) + extinit!(gv, false) + constant!(gv, true) + unnamed_addr!(gv, true) + end + end + empty!(relocs.records) + return +end + +""" + emit_patchable_relocations!(mod, relocs) + +Emit slots as writable, null-initialized definitions, and leave interior records as the +`extinit` definitions they already are. The loader must patch every record by `(name, +offset)` after loading the object ([`resolved_relocations`](@ref)). +""" +function emit_patchable_relocations!(mod::LLVM.Module, relocs::Relocations) + used = GlobalVariable[] + foreach_relocation(mod, relocs) do rec, gv + if rec.kind === SlotSite + initializer!(gv, null(global_value_type(gv))) + constant!(gv, false) + extinit!(gv, true) + end + # Two objects can define the same record: a relocation-carrying runtime-library + # function keeps its own job's namespace in every kernel it is linked into. A loader + # holding both in one symbol namespace (an ORC `JITDylib`) would see a duplicate + # definition, so define them weakly and let it coalesce. That is sound rather than + # merely quiet: a shared name means a shared producing job, hence the same target + # (`add_relocation!` enforces agreement), so whichever definition survives gets + # patched with the word every object referencing it expects. `llvm.used` below still + # anchors them against DCE, and `externally_initialized` still stops the optimizer + # from believing the null initializer. + linkage!(gv, LLVM.API.LLVMWeakODRLinkage) + push!(used, gv) + end + isempty(used) || set_used!(mod, used...) + return +end + +# The functions whose bodies use `value`, following constant expressions (a `getelementptr` +# onto a global, an isbits union's `{ptr, i8}` aggregate) through to the instructions they end +# up in. +function using_functions!(fns::Set{LLVM.Function}, @nospecialize(value)) + for use in uses(value) + val = user(use) + if val isa LLVM.Instruction + push!(fns, LLVM.parent(LLVM.parent(val))) + elseif val isa LLVM.Constant + using_functions!(fns, val) + end + end + return fns +end + +# A relocation word comes out of a table whose base the back-end derives from per-dispatch +# state, which only an entry point can reach (a kernel's state argument, typically). So hoist +# every *other* function still holding a relocation use into its caller(s): mark it +# `alwaysinline` and run the inliner, until only entry points hold one. Mirrors +# `inline_unreachable_control_flow!`, and handles the `entry → A → B` case for the same reason: +# `A` gets marked on the next round, once `B` has been inlined into it. +function inline_relocation_users!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + alwaysinline_attr = EnumAttribute("alwaysinline", 0) + noinline_attr = EnumAttribute("noinline", 0) + # by name: a function we mark may well be gone by the next round + hoisted = Set{String}() + while true + users = Set{LLVM.Function}() + for rec in relocs.records + haskey(globals(mod), rec.name) || continue + using_functions!(users, globals(mod)[rec.name]) + end + + marked = false + for f in users + # an entry point has no call sites, and is where the state arrives + isempty(uses(f)) && continue + fn = LLVM.name(f) + fn in hoisted && + error("""Function `$fn` uses a relocation but could not be inlined into an + entry point (it is likely recursive or address-taken), so it cannot + reach the relocation table.""") + push!(hoisted, fn) + attrs = function_attributes(f) + delete!(attrs, noinline_attr) + push!(attrs, alwaysinline_attr) + marked = true + end + marked || break + + @dispose pb=NewPMPassBuilder() begin + add!(pb, AlwaysInlinerPass()) + run!(pb, mod, llvm_machine(job.config.target)) + end + end + return +end + +""" + emit_table_relocations!(job, mod, relocs) + +Rewrite every record into an indexed load from a back-end-provided table of words, the +`:table` strategy's lowering. A record's index is its rank in `relocs`, which the lowering +copies into `relocs.table` so that [`resolved_relocation_table`](@ref) delivers the words in +that same order regardless of what happens to the records afterwards. + +Slots become `load(gep(base, index))` and are erased. Interior boxes cannot be patched +after load — the platforms needing this have no writable program-scope storage — so each is +demoted to a per-function stack copy whose header word comes from the table. +[`relocation_table_pointer`](@ref) supplies the base pointer; since it can only do so where +the state is available, callees still holding a relocation use are inlined first. +""" +function emit_table_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + isempty(relocs) && return + LLVM.version() >= v"17" || + error("The `:table` relocation lowering requires LLVM 17 or later (Julia 1.12+)") + + # Fix the word order up front, in its own vector: the indices below are baked into the + # code, so what the loader delivers must not depend on the manifest still being in this + # order afterwards. + empty!(relocs.table) + append!(relocs.table, (rec.target for rec in relocs.records)) + + inline_relocation_users!(job, mod, relocs) + + T_word = relocation_word_type() + + # One base pointer per function, materialized at the top of its entry block (the state + # it derives from is a function argument, so it dominates every use). + bases = Dict{LLVM.Function, LLVM.Value}() + function table_base(f::LLVM.Function) + get!(bases, f) do + @dispose builder=IRBuilder() begin + position!(builder, first(instructions(first(blocks(f))))) + relocation_table_pointer(job, builder, f) + end + end + end + function table_word(builder::IRBuilder, index::Int) + f = LLVM.parent(position(builder)) + ptr = inbounds_gep!(builder, T_word, table_base(f), + [ConstantInt(LLVM.Int32Type(), index - 1)]) + load!(builder, T_word, ptr) + end + + mod_gvs = globals(mod) + for (index, rec) in enumerate(relocs.records) + haskey(mod_gvs, rec.name) || error("Missing relocation global '$(rec.name)'") + gv = mod_gvs[rec.name] + check_relocation(mod, rec, gv) + + if rec.kind === SlotSite + rewrite_word_loads!(gv, "relocation slot '$(rec.name)'") do builder + table_word(builder, index) + end + prune_constexpr_uses!(gv) + isempty(uses(gv)) || + error("Relocation slot '$(rec.name)' still has uses after redirection") + erase!(gv) + else + demote_relocatable_box!(mod, gv, rec, table_word, index) + end + end + return +end + +# Copy a relocatable box into a per-function stack slot and fill its header from the +# relocation table. Sound because a box address carries no identity of its own: `isbits` egal +# compares by content, so a per-invocation copy is indistinguishable from a shared one. +function demote_relocatable_box!(mod::LLVM.Module, gv::GlobalVariable, rec::Relocation, + table_word, index::Int) + boxty = global_value_type(gv)::LLVM.StructType + init = initializer(gv) + header_idx = Int(element_at(datalayout(mod), boxty, rec.offset)) + + allocas = Dict{LLVM.Function, LLVM.Value}() + function box_alloca(f::LLVM.Function) + get!(allocas, f) do + @dispose builder=IRBuilder() begin + position!(builder, first(instructions(first(blocks(f))))) + ptr = alloca!(builder, boxty) + # keep Julia's heap alignment, which the payload's `isbits` layout assumes + alignment!(ptr, max(alignment(gv), 16)) + store!(builder, init, ptr) + # overwrite the (zeroed) header field with the resolved relocation word + word = table_word(builder, index) + store!(builder, word, struct_gep!(builder, boxty, ptr, header_idx)) + ptr + end + end + end + replace_global_with_local!(gv, box_alloca) + return +end + +""" + apply_relocations!(mod, relocs) + +Resolve every live record into `mod` without consuming `relocs`, so cached metadata can be +reused. Records whose global was optimized away are skipped. Resolution permanently roots +referenced Julia values in the process. Apply once per parsed module. + +For consumers that need a session-resolved copy of a module whose cached form is symbolic — +e.g. to read a type tag out of the IR — alongside the symbolic one they cache. +""" +function apply_relocations!(mod::LLVM.Module, relocs::Relocations) + live = copy(relocs) + prune_dead_relocations!(mod, live) + bake_relocations!(mod, live) + return +end + + +## introspection + +function referenced_object(value, relocs::Relocations) + # This is best-effort: optimized shapes fall back to the unknown-binding error path. + while value isa ConstantExpr && + opcode(value) in (LLVM.API.LLVMBitCast, LLVM.API.LLVMAddrSpaceCast) + value = first(operands(value)) + end + if value isa LLVM.LoadInst + source = first(operands(value)) + while source isa ConstantExpr && + opcode(source) in (LLVM.API.LLVMBitCast, LLVM.API.LLVMAddrSpaceCast) + source = first(operands(source)) + end + if source isa GlobalVariable + rec = find_relocation(relocs, LLVM.name(source)) + if rec !== nothing && rec.target isa JuliaValueRef + return Some(rec.target.value) + end + end + elseif value isa ConstantExpr && opcode(value) == LLVM.API.LLVMIntToPtr + ptr = Ptr{Cvoid}(convert(Int, first(operands(value)))) + return Some(Base.unsafe_pointer_to_objref(ptr)) + end + return nothing +end diff --git a/src/rtlib.jl b/src/rtlib.jl index b60ba806..d7bdd87b 100644 --- a/src/rtlib.jl +++ b/src/rtlib.jl @@ -60,22 +60,20 @@ end ## functionality to build the runtime library -# Per-function compilation results for the GPU runtime library, cached through the -# same `cached_results` mechanism back-ends use for kernels. On 1.11+ the bitcode -# thus lives on the runtime function's `CodeInstance` — possibly alongside a -# back-end's own results struct — and persists through precompilation, so sessions -# loading a back-end that compiled its runtime during precompile skip codegen -# entirely. On 1.10 it is cached for the duration of the session. +# Per-function relocatable bitcode for the GPU runtime library. When Julia exposes the +# required relocation metadata, this persists with the function's `CodeInstance`. +# `runtime_libs` below provides the session cache on older Julia versions. mutable struct RuntimeFunctionResults bitcode::Union{Nothing,Vector{UInt8}} - RuntimeFunctionResults() = new(nothing) + relocations::Relocations + RuntimeFunctionResults() = new(nothing, Relocations()) end # Compile a single runtime function and link it into `mod`. The renamed bitcode is # memoized through `RuntimeFunctionResults`; the session-local `runtime_libs` cache # below additionally avoids repeating the parse-and-link work within a session. -function emit_function!(mod, config::CompilerConfig, source::MethodInstance, method, - world::UInt) +function emit_function!(mod, relocs::Relocations, config::CompilerConfig, + source::MethodInstance, method, world::UInt) name = method.llvm_name rt_job = CompilerJob(source, config, world) @@ -83,12 +81,16 @@ function emit_function!(mod, config::CompilerConfig, source::MethodInstance, met # inference itself. ci, res = runtime_function_results(rt_job) if res !== nothing && res.bitcode !== nothing - link!(mod, parse(LLVM.Module, MemoryBuffer(res.bitcode))) + link_relocatable!(mod, relocs, + parse(LLVM.Module, MemoryBuffer(res.bitcode)), + res.relocations) ci === nothing && (ci = runtime_code_instance(rt_job)) return ci::CodeInstance end - new_mod, meta = compile_unhooked(:llvm, rt_job) + # Keep this intermediate module relocatable even when the final back-end resolves + # relocations eagerly. The caller links a fresh copy and lowers the merged sites. + new_mod, meta = compile_unhooked(:llvm, rt_job; resolve_relocations=false) ft = function_type(meta.entry) expected_ft = convert(LLVM.FunctionType, method) if return_type(ft) != return_type(expected_ft) @@ -97,13 +99,7 @@ function emit_function!(mod, config::CompilerConfig, source::MethodInstance, met # recent Julia versions include prototypes for all runtime functions, even if unused run!(StripDeadPrototypesPass(), new_mod, llvm_machine(config.target)) - - # Resolve constgv mappings before their metadata is discarded. Dedicated Bool - # globals are resolved after linking into the toplevel module. - if !isempty(meta.gv_to_value) - portable = relocate_gvs!(new_mod, meta.gv_to_value) - portable || mark_session_dependent!(rt_job) - end + prune_dead_relocations!(new_mod, meta.relocations) # rename to the final `gpu_*` name on the per-function module, so the cached bitcode # is immediately link-ready (no per-session rename pass on a cache hit). @@ -118,17 +114,29 @@ function emit_function!(mod, config::CompilerConfig, source::MethodInstance, met io = IOBuffer() write(io, new_mod) ci === nothing && (ci = runtime_code_instance(rt_job)) - res === nothing && (res = job_results(RuntimeFunctionResults, ci, rt_job.config)) - res.bitcode = take!(io) + if supports_relocatable_ir() + res === nothing && (res = runtime_results(RuntimeFunctionResults, ci, rt_job.config)) + res.bitcode = take!(io) + res.relocations = meta.relocations + end - link!(mod, new_mod) + link_relocatable!(mod, relocs, new_mod, meta.relocations) return ci::CodeInstance end +function runtime_results(::Type{V}, ci::CodeInstance, config::CompilerConfig) where {V} + @static if HAS_INTEGRATED_CACHE + if supports_relocatable_ir() + return persistent_results(V, ci, config) + end + end + return session_results(V, ci, config) +end + function runtime_function_results(@nospecialize(job::CompilerJob)) ci = job_code_instance(job) ci === nothing && return nothing, nothing - return ci, job_results(RuntimeFunctionResults, ci, job.config) + return ci, runtime_results(RuntimeFunctionResults, ci, job.config) end function runtime_method_instance(@nospecialize(job::CompilerJob), method) @@ -170,13 +178,14 @@ function build_runtime(@nospecialize(job::CompilerJob), config::CompilerConfig) mod = LLVM.Module("GPUCompiler run-time library") sources = MethodInstance[] code_instances = CodeInstance[] + relocs = Relocations() for method in values(Runtime.methods) resolved = runtime_method_instance(job, method) resolved === nothing && continue source = resolved push!(sources, source) - push!(code_instances, emit_function!(mod, config, source, method, job.world)) + push!(code_instances, emit_function!(mod, relocs, config, source, method, job.world)) end # we cannot optimize the runtime library, because the code would then be optimized again @@ -184,13 +193,13 @@ function build_runtime(@nospecialize(job::CompilerJob), config::CompilerConfig) # removes Julia address spaces, which would then lead to type mismatches when using # functions from the runtime library from IR that has not been stripped of AS info. - return mod, sources, code_instances + return mod, sources, code_instances, relocs end # Session-local cache of assembled runtime libraries, keyed by # `(runtime_config(job), opaque_pointers)`: the derived runtime config covers every -# codegen-relevant setting (e.g. the debug level, which is baked into the runtime IR -# as a constant), while cosmetic kernel-job fields are normalized away. Cross-session +# codegen-relevant setting (e.g. the debug level stored in the runtime IR), while +# cosmetic kernel-job fields are normalized away. Cross-session # persistence happens at the per-function level (see `RuntimeFunctionResults`): # reassemble on first use of each session, then reuse within the session. # @@ -203,6 +212,7 @@ mutable struct RuntimeLibrary sources::Vector{MethodInstance} code_instances::Vector{CodeInstance} validated_world::UInt + relocations::Relocations end function runtime_library_valid(lib::RuntimeLibrary, @nospecialize(job::CompilerJob)) @@ -235,14 +245,16 @@ const runtime_libs_lock = ReentrantLock() cached = Base.@lock runtime_libs_lock begin cached = get(runtime_libs, key, nothing) if cached === nothing || !runtime_library_valid(cached, job) - lib, sources, code_instances = build_runtime(job, config) + lib, sources, code_instances, relocations = build_runtime(job, config) io = IOBuffer() write(io, lib) - cached = RuntimeLibrary(take!(io), sources, code_instances, job.world) + cached = RuntimeLibrary(take!(io), sources, code_instances, job.world, + relocations) runtime_libs[key] = cached end cached end - return parse(LLVM.Module, MemoryBuffer(cached.bytes); lazy=true) + return parse(LLVM.Module, MemoryBuffer(cached.bytes); lazy=true), + cached.relocations end diff --git a/src/utils.jl b/src/utils.jl index db2fc7d7..45f13045 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -274,6 +274,33 @@ function prune_constexpr_uses!(root::LLVM.Value) end end +## replacing a global with a runtime value + +# Replace every use of `gv` with the function-local value `replacement(f)`, then erase it. +# Plain `replace_uses!` can't do this: an instruction is not a valid operand of the constant +# expressions/aggregates the global's address may be folded into (a `getelementptr` onto it, +# an isbits union's `{ptr, i8}` return value), so those constants are expanded into +# instructions first (phi operands materialize in their incoming block). `replacement` is +# invoked once per using function (memoize per-function state such as an entry-block alloca +# in the callback). +function replace_global_with_local!(gv::LLVM.GlobalVariable, replacement) + convert_users_to_instructions!([gv]) + for use in collect(uses(gv)) + inst = user(use) + inst isa LLVM.Instruction || + error("Unexpected use of global '$(LLVM.name(gv))': $inst") + f = LLVM.parent(LLVM.parent(inst)) + ops = operands(inst) + for i in 1:length(ops) + ops[i] == gv || continue + ops[i] = replacement(f) + end + end + @assert isempty(uses(gv)) "global '$(LLVM.name(gv))' still has uses after replacement" + erase!(gv) + return +end + ## function-signature rewriting @@ -418,3 +445,14 @@ end return inits end end + +"""Whether Julia exposes enough global-variable metadata to emit relocatable IR.""" +supports_relocatable_ir() = @static if VERSION >= v"1.13.0-DEV.623" + true +else + # `jl_get_llvm_gvs_globals` was backported to 1.10, so the symbol alone is not enough: + # 1.10's codegen still embeds Julia addresses (as `inttoptr` constants) in the JIT + # (non-imaging) mode we compile in, instead of emitting the relocatable global + # declarations the relocation machinery collects. Only 1.11+ emits those declarations. + VERSION >= v"1.11-" && HAS_LLVM_GVS_GLOBALS +end diff --git a/src/validation.jl b/src/validation.jl index c8e97638..0e19d440 100644 --- a/src/validation.jl +++ b/src/validation.jl @@ -162,8 +162,8 @@ end # `show` via `showerror`, avoiding the default field-dump that derefs disposed IR Base.show(io::IO, err::InvalidIRError) = showerror(io, err) -function check_ir(job, args...) - errors = check_ir!(job, IRError[], args...) +function check_ir(job, mod::LLVM.Module, relocs::Relocations=Relocations()) + errors = check_ir!(job, IRError[], mod, relocs) unique!(errors) if !isempty(errors) throw(InvalidIRError(job, errors)) @@ -172,9 +172,9 @@ function check_ir(job, args...) return end -function check_ir!(job, errors::Vector{IRError}, mod::LLVM.Module) +function check_ir!(job, errors::Vector{IRError}, mod::LLVM.Module, relocs::Relocations) for f in functions(mod) - check_ir!(job, errors, f) + check_ir!(job, errors, f, relocs) end # custom validation @@ -183,10 +183,10 @@ function check_ir!(job, errors::Vector{IRError}, mod::LLVM.Module) return errors end -function check_ir!(job, errors::Vector{IRError}, f::LLVM.Function) +function check_ir!(job, errors::Vector{IRError}, f::LLVM.Function, relocs::Relocations) for bb in blocks(f), inst in instructions(bb) if isa(inst, LLVM.CallInst) - check_ir!(job, errors, inst) + check_ir!(job, errors, inst, relocs) elseif isa(inst, LLVM.LoadInst) check_ir!(job, errors, inst) end @@ -221,7 +221,7 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.LoadInst) return errors end -function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst) +function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst, relocs::Relocations) bt = backtrace(inst) dest = called_operand(inst) if isa(dest, LLVM.Function) @@ -233,11 +233,9 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst) elseif fn == "jl_get_binding_or_error" || fn == "ijl_get_binding_or_error" try m, sym = arguments(inst) - sym = first(operands(sym::ConstantExpr))::ConstantInt - sym = convert(Int, sym) - sym = Ptr{Cvoid}(sym) - sym = Base.unsafe_pointer_to_objref(sym) - push!(errors, (DELAYED_BINDING, bt, sym)) + ref = referenced_object(sym, relocs) + ref === nothing && error("Unknown binding") + push!(errors, (DELAYED_BINDING, bt, something(ref))) catch e @safe_debug "Decoding arguments to jl_get_binding_or_error failed" inst bb=LLVM.parent(inst) push!(errors, (DELAYED_BINDING, bt, nothing)) @@ -246,10 +244,9 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst) fn == "jl_get_binding_value_seqcst" || fn == "ijl_get_binding_value_seqcst" try # pry the binding from the IR - expr = arguments(inst)[1]::ConstantExpr - expr = first(operands(expr))::ConstantInt # get rid of inttoptr - ptr = Ptr{Any}(convert(Int, expr)) - obj = Base.unsafe_pointer_to_objref(ptr) + ref = referenced_object(arguments(inst)[1], relocs) + ref === nothing && error("Unknown binding") + obj = something(ref) push!(errors, (DELAYED_BINDING, bt, obj.globalref)) catch e @safe_debug "Decoding arguments to jl_reresolve_binding_value_seqcst failed" inst bb=LLVM.parent(inst) @@ -258,10 +255,9 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst) elseif startswith(fn, "tojlinvoke") try fun, args, nargs = arguments(inst) - fun = first(operands(fun::ConstantExpr))::ConstantInt - fun = convert(Int, fun) - fun = Ptr{Cvoid}(fun) - fun = Base.unsafe_pointer_to_objref(fun)::Base.Function + ref = referenced_object(fun, relocs) + ref === nothing && error("Unknown function") + fun = something(ref)::Base.Function push!(errors, (DYNAMIC_CALL, bt, fun)) # XXX: an invoke trampoline happens when codegen doesn't have access to code # which suggests a GPUCompiler.jl bug. throw an error instead? @@ -279,10 +275,9 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst) end try fun, args, nargs, meth = arguments(inst) - meth = first(operands(meth::ConstantExpr))::ConstantInt - meth = convert(Int, meth) - meth = Ptr{Cvoid}(meth) - meth = Base.unsafe_pointer_to_objref(meth)::Core.MethodInstance + ref = referenced_object(meth, relocs) + ref === nothing && error("Unknown method instance") + meth = something(ref)::Core.MethodInstance push!(errors, (DYNAMIC_CALL, bt, meth.def)) catch e @safe_debug "Decoding arguments to jl_invoke failed" inst bb=LLVM.parent(inst) @@ -291,10 +286,9 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst) elseif fn == "jl_apply_generic" || fn == "ijl_apply_generic" try f, args, nargs = arguments(inst) - f = first(operands(f))::ConstantInt # get rid of inttoptr - f = convert(Int, f) - f = Ptr{Cvoid}(f) - f = Base.unsafe_pointer_to_objref(f) + ref = referenced_object(f, relocs) + ref === nothing && error("Unknown function") + f = something(ref) push!(errors, (DYNAMIC_CALL, bt, f)) catch e @safe_debug "Decoding arguments to jl_apply_generic failed" inst bb=LLVM.parent(inst) diff --git a/test/helpers/metal.jl b/test/helpers/metal.jl index 8bea2a25..5346bf78 100644 --- a/test/helpers/metal.jl +++ b/test/helpers/metal.jl @@ -26,6 +26,37 @@ end ThreadedRuntimeCompilerJob = CompilerJob{MetalCompilerTarget,ThreadedRuntimeCompilerParams} GPUCompiler.runtime_module(::ThreadedRuntimeCompilerJob) = ThreadedRuntime +# Selects the `:table` relocation strategy, as Metal.jl does, so relocation words are read out +# of a table reached through the kernel state instead of being baked. Used to test that delivery +# path (and the Metal target's `relocation_table_pointer`) without depending on Metal.jl. +struct TableKernelState + reloc_table::Core.LLVMPtr{UInt64, 1} +end + +struct TableCompilerParams <: AbstractCompilerParams end +TableCompilerJob = CompilerJob{MetalCompilerTarget,TableCompilerParams} +GPUCompiler.runtime_module(::TableCompilerJob) = TestRuntime +# as Metal.jl does: only a kernel has the state a table can be reached through, so reflection +# on a plain device function falls back to session-local resolution +GPUCompiler.relocation_lowering(job::TableCompilerJob) = + job.config.kernel ? (:table) : (:bake) +GPUCompiler.kernel_state_type(::TableCompilerJob) = TableKernelState + +function create_table_job(@nospecialize(func), @nospecialize(types); kwargs...) + config_kwargs, kwargs = split_kwargs(kwargs, GPUCompiler.CONFIG_KWARGS) + source = methodinstance(typeof(func), Base.to_tuple_type(types), Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"12.2", metal=v"3.0", air=v"3.0") + config = CompilerConfig(target, TableCompilerParams(); kernel=false, config_kwargs...) + CompilerJob(source, config), kwargs +end + +function code_native_table(io::IO, @nospecialize(func), @nospecialize(types); kwargs...) + job, kwargs = create_table_job(func, types; kwargs...) + GPUCompiler.code_native(io, job; kwargs...) +end +code_native_table(@nospecialize(func), @nospecialize(types); kwargs...) = + code_native_table(stdout, func, types; kwargs...) + function create_job(@nospecialize(func), @nospecialize(types); kwargs...) config_kwargs, kwargs = split_kwargs(kwargs, GPUCompiler.CONFIG_KWARGS) source = methodinstance(typeof(func), Base.to_tuple_type(types), Base.get_world_counter()) diff --git a/test/helpers/native.jl b/test/helpers/native.jl index 63d8a0f3..5e63e992 100644 --- a/test/helpers/native.jl +++ b/test/helpers/native.jl @@ -1,6 +1,7 @@ module Native using ..GPUCompiler +using LLVM import ..TestRuntime # local method table for device functions @@ -9,9 +10,11 @@ Base.Experimental.@MethodTable(test_method_table) struct CompilerParams <: AbstractCompilerParams entry_safepoint::Bool method_table + relocations::Symbol - CompilerParams(entry_safepoint::Bool=false, method_table=test_method_table) = - new(entry_safepoint, method_table) + CompilerParams(entry_safepoint::Bool=false, method_table=test_method_table, + relocations::Symbol=:bake) = + new(entry_safepoint, method_table, relocations) end module Runtime end @@ -22,16 +25,97 @@ GPUCompiler.runtime_module(::NativeCompilerJob) = Runtime GPUCompiler.method_table(@nospecialize(job::NativeCompilerJob)) = job.config.params.method_table GPUCompiler.can_safepoint(@nospecialize(job::NativeCompilerJob)) = job.config.params.entry_safepoint +# Every mode ends up in an ORC JIT: `patch` emits definitions for `load` to write after adding +# the object, while `table` delivers the words as run-time data, reached through the single +# patchable global below. The latter is a stand-in for a platform that offers no access to +# loaded code at all (Metal), letting that strategy be tested off-device. +GPUCompiler.relocation_lowering(@nospecialize(job::NativeCompilerJob)) = + job.config.params.relocations + +# The `:table` back-end contract: hand out a pointer to the table base. A real back-end reads +# it out of per-dispatch state; here one patchable global holds it for the whole object, which +# `load` fills in after adding the object to the JIT. +const RELOC_TABLE_BASE = "__reloc_table_base" + +function GPUCompiler.relocation_table_pointer(@nospecialize(job::NativeCompilerJob), + builder::LLVM.IRBuilder, fun::LLVM.Function) + mod = LLVM.parent(fun) + T_word = GPUCompiler.relocation_word_type() + gv = if haskey(globals(mod), RELOC_TABLE_BASE) + globals(mod)[RELOC_TABLE_BASE] + else + gv = GlobalVariable(mod, T_word, RELOC_TABLE_BASE) + initializer!(gv, LLVM.ConstantInt(T_word, 0)) + extinit!(gv, true) + linkage!(gv, LLVM.API.LLVMExternalLinkage) + set_used!(mod, gv) + gv + end + return inttoptr!(builder, load!(builder, T_word, gv), LLVM.PointerType(T_word)) +end + +function GPUCompiler.mcgen(@nospecialize(job::NativeCompilerJob), mod::LLVM.Module, + format=LLVM.API.LLVMAssemblyFile) + if job.config.params.relocations !== :bake + target = job.config.target + @dispose tm=JITTargetMachine(GPUCompiler.llvm_triple(target), target.cpu, + target.features) begin + return String(emit(tm, mod, format)) + end + else + return invoke(GPUCompiler.mcgen, Tuple{CompilerJob,LLVM.Module,Any}, + job, mod, format) + end +end + function create_job(@nospecialize(func), @nospecialize(types); - entry_safepoint::Bool=false, method_table=test_method_table, kwargs...) + entry_safepoint::Bool=false, method_table=test_method_table, + relocations::Symbol=:bake, kwargs...) config_kwargs, kwargs = split_kwargs(kwargs, GPUCompiler.CONFIG_KWARGS) source = methodinstance(typeof(func), Base.to_tuple_type(types), Base.get_world_counter()) target = NativeCompilerTarget(;jlruntime=true) - params = CompilerParams(entry_safepoint, method_table) + params = CompilerParams(entry_safepoint, method_table, relocations) config = CompilerConfig(target, params; kernel=false, config_kwargs...) CompilerJob(source, config), kwargs end +# Add an object to a fresh ORC JIT and supply its relocation words, the loader in miniature. +# Under `:patch` each word is written into the loaded image by name and offset (CUDA does the +# same with `cuModuleGetGlobal` + `cuMemcpyHtoD`); under `:table` one word table is allocated +# and its address written into the module's single table-base global (Metal passes the same +# address in the kernel state). Pass empty `relocs` for objects that need neither. +# +# The returned table must stay rooted for as long as the code is callable. +function load(obj::Vector{UInt8}, entry::String, relocs::GPUCompiler.Relocations; + table::Bool=false) + lljit = LLJIT(; tm=JITTargetMachine()) + try + jd = JITDylib(lljit) + prefix = LLVM.get_prefix(lljit) + add!(jd, LLVM.CreateDynamicLibrarySearchGeneratorForProcess(prefix)) + + add!(lljit, jd, MemoryBuffer(obj)) + words = UInt[] + if table + words = GPUCompiler.resolved_relocation_table(relocs) + if !isempty(words) + base = lookup(lljit, RELOC_TABLE_BASE) + unsafe_store!(Ptr{UInt}(pointer(base)), UInt(pointer(words))) + end + else + for (rec, value) in GPUCompiler.resolved_relocations(relocs) + addr = lookup(lljit, rec.name) + unsafe_store!(Ptr{UInt}(pointer(addr) + rec.offset), value) + end + end + addr = lookup(lljit, entry) + return pointer(addr), lljit, words + catch + dispose(lljit) + rethrow() + end +end + function code_typed(@nospecialize(func), @nospecialize(types); kwargs...) job, kwargs = create_job(func, types; kwargs...) GPUCompiler.code_typed(job; kwargs...) diff --git a/test/helpers/ptx.jl b/test/helpers/ptx.jl index 4b00f707..d2233325 100644 --- a/test/helpers/ptx.jl +++ b/test/helpers/ptx.jl @@ -3,10 +3,17 @@ module PTX using ..GPUCompiler import ..TestRuntime -struct CompilerParams <: AbstractCompilerParams end +struct CompilerParams <: AbstractCompilerParams + patch::Bool + CompilerParams(patch::Bool=false) = new(patch) +end PTXCompilerJob = CompilerJob{PTXCompilerTarget,CompilerParams} +# `patch=true` keeps relocations symbolic (as CUDA.jl does); plain jobs resolve them in IR. +GPUCompiler.relocation_lowering(@nospecialize(job::PTXCompilerJob)) = + job.config.params.patch ? :patch : :bake + struct PTXKernelState data::Int64 end @@ -39,14 +46,14 @@ function create_job(@nospecialize(func), @nospecialize(types); cap=v"7.0", ptx=v"6.0", feature_set=:baseline, minthreads=nothing, maxthreads=nothing, blocks_per_sm=nothing, maxregs=nothing, - fastmath=false, + fastmath=false, patch::Bool=false, kwargs...) config_kwargs, kwargs = split_kwargs(kwargs, GPUCompiler.CONFIG_KWARGS) source = methodinstance(typeof(func), Base.to_tuple_type(types), Base.get_world_counter()) target = PTXCompilerTarget(; cap, ptx, feature_set, minthreads, maxthreads, blocks_per_sm, maxregs, fastmath) - params = CompilerParams() + params = CompilerParams(patch) config = CompilerConfig(target, params; kernel=false, config_kwargs...) CompilerJob(source, config), kwargs end diff --git a/test/metal.jl b/test/metal.jl index 763b9376..6546940d 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -388,11 +388,200 @@ end Tuple{Core.LLVMPtr{Int32,1}, Bool, Int32}; dump_module=true, kernel=true) end - @test occursin("@jl_true_box = private unnamed_addr addrspace(2) constant", ir) + # Either codegen references `jl_true` as a named global and we materialize it as a + # device-resident box, or (newer 1.14-DEV) codegen constructs the smalltag box on + # the stack itself and no `jl_true` reference exists at all. Both satisfy the + # invariant this testset guards: no unresolved `jl_true` reaches the back-end. + @test occursin("@jl_true_box = private unnamed_addr addrspace(2) constant", ir) || + !occursin("jl_true", ir) @test !occursin("@jl_true = external", ir) end end +@testset "relocations through the kernel state" begin + # Under the `:table` strategy (Metal.jl's choice), GPUCompiler rewrites each relocation + # into an indexed load from a word table the loader passes in the kernel state, rather + # than baking a session address into the module. Compile a relocation-carrying kernel and + # check the emitted AIR reads its words that way (no device needed — this runs through + # the LLVM downgrader). + if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17" + mod = @eval module $(gensym()) + function kernel(ptr::Core.LLVMPtr{UInt,1}) + Base.unsafe_store!(ptr, UInt(pointer_from_objref(:table_probe))) + return + end + end + + # the relocation records the kernel carries (at the `:llvm` level, before lowering) + job, _ = Metal.create_table_job(mod.kernel, (Core.LLVMPtr{UInt,1},); kernel=true) + relocs = JuliaContext() do ctx + _, meta = GPUCompiler.compile_unhooked(:llvm, job; resolve_relocations=false) + meta.relocations + end + @test !isempty(relocs) + @test any(rec -> rec.target isa GPUCompiler.JuliaValueRef, relocs.records) + + air = sprint() do io + Metal.code_native_table(io, mod.kernel, (Core.LLVMPtr{UInt,1},); kernel=true) + end + # the table base is loaded out of the kernel-state argument, and the words out of the + # table -- a bake would instead leave a private constant holding the resolved address + @test occursin("reloc_table", air) + @test occursin(r"load i64, i64 addrspace\(1\)\*", air) + # nothing is left of the site globals the records named + for rec in relocs.records + @test !occursin("@$(rec.name) ", air) + end + end +end + +@testset "codegen counter normalization" begin + # Julia's per-session codegen counter has to be scrubbed from everything that reaches the + # bitcode, or the metallib is not reproducible: symbol names, the block labels inlining + # synthesizes from them, and the metadata strings (orphaned `DISubprogram` linkage names, + # Julia's alias scopes) that mention them. The rewrite is deterministic but it must not + # reach into names that merely resemble the pattern. + JuliaContext() do ctx + m = parse(LLVM.Module, """ + define void @julia_probe_4242() { + julia_inlinee_77.exit: + ret void, !alias.scope !0 + } + + define void @myjulia_probe_4242() { + ret void + } + + define void @julia_probe_12bar() { + ret void + } + + define void @"julia_record_exception!_18521"() { + ret void, !alias.scope !2 + } + + define void @"julia_#closure#42_777"() { + ret void + } + + !0 = !{!1} + !1 = distinct !{!1, !"julia_inlinee_77: %union_bytes_return"} + !2 = !{!3} + !3 = distinct !{!3, !"julia_record_exception!_18521"}""") + GPUCompiler.normalize_julia_symbol_names!(m) + ir = string(m) + + # a codegen name is replaced by a deterministic module-local rank, wherever it sits + @test !occursin("@julia_probe_4242", ir) + @test !occursin("julia_inlinee_77", ir) + @test occursin("@julia_probe_1(", ir) + @test occursin(r"julia_inlinee_[0-9]+\.exit", ir) + @test occursin(r"!\"julia_inlinee_[0-9]+: %union_bytes_return\"", ir) + + # names outside `\w` — mutating `!`, closure `#` — are codegen names too + @test !occursin("julia_record_exception!_18521", ir) + @test !occursin("julia_#closure#42_777", ir) + @test occursin(r"@\"julia_record_exception!_[0-9]+\"", ir) + @test occursin(r"!\"julia_record_exception!_[0-9]+\"", ir) + @test occursin(r"@\"julia_#closure#42_[0-9]+\"", ir) + + # ...but a user symbol that merely ends or contains the pattern is left alone + @test occursin("@myjulia_probe_4242", ir) + @test occursin("@julia_probe_12bar", ir) + dispose(m) + end +end + +@testset "identity kernel arguments" begin + # A boxed argument that gets past `check_invocation` has no fields, so it can only be used + # by identity (`sym === :foo`, whose comparison target is itself a relocation). Metal cannot + # pass a heap reference, so `add_parameter_address_spaces!` lowers the parameter to a buffer + # holding the object's bare address word — which the host supplies — and the body turns that + # word back into the pointer Julia's comparison expects. + mod = @eval module $(gensym()) + function kernel(ptr, sym::Symbol) + unsafe_store!(ptr, sym === :identity_arg_probe ? 1 : 2) + return + end + end + tt = (Core.LLVMPtr{Int,1}, Symbol) + + ir = sprint(io -> Metal.code_llvm(io, mod.kernel, tt; kernel=true)) + # the body loads the word and compares it against the (here: baked) target. On LLVM 20+ + # the comparison against the pointer constant folds down to the bare words (`icmp eq + # i64`); older LLVM keeps the pointer-typed compare. Both mean the same thing. + @test occursin("load i64", ir) + @test occursin(r"icmp eq (i64|ptr|\{\}\*)", ir) + # nothing is left in a Julia address space, which AIR has no equivalent for + @test !occursin("addrspace(10)", ir) + + # A `Symbol` has no size of its own, so the AIR metadata must describe the address word. + air = sprint(io -> Metal.code_native(io, mod.kernel, tt; kernel=true)) + sym_arg = only(filter(l -> occursin("!\"air.arg_name\", !\"sym\"", l), split(air, '\n'))) + @test occursin("!\"air.arg_type_name\", !\"Symbol\"", sym_arg) + @test occursin("!\"air.arg_type_size\", i32 $(sizeof(UInt))", sym_arg) + @test occursin("!\"air.arg_type_align_size\", i32 $(sizeof(UInt))", sym_arg) + # read-only: unlike a data pointer, the slot is never written back through + @test occursin("!\"air.read\"", sym_arg) +end + +@testset "boxed constant classification" begin + JuliaContext() do ctx + mod = LLVM.Module("boxed-constant-classification") + T = LLVM.StructType([LLVM.Int64Type(), LLVM.Int64Type()]) + + function private_constant(name) + gv = GlobalVariable(mod, T, name) + initializer!(gv, ConstantStruct(LLVM.Constant[ConstantInt(0), ConstantInt(0)])) + linkage!(gv, LLVM.API.LLVMPrivateLinkage) + constant!(gv, true) + unnamed_addr!(gv, true) + return gv + end + + @test GPUCompiler.is_boxed_constant(private_constant("value_box")) + @test !GPUCompiler.is_boxed_constant(private_constant("ordinary_constant")) + end +end + +@testset "smalltag boxed union constant demotion" begin + # A non-relocatable, smalltag-header isbits-`Union` constant is materialized as a private + # constant box. When its payload address escapes the `{ptr, i8}` union return, the box must + # be demoted to a stack `alloca` (`demote_boxed_constants!`), or `add_global_address_spaces!` + # sinks it into addrspace(2), whose payload the kernel then misreads at run time (AIR has + # no generic address space, so the addrspacecast to AS 0 folded into the returned aggregate + # makes the load read thread memory). This is independent of relocations — the kernel + # carries none — and regresses easily since neither the native tests (a different target) nor + # the other metal tests exercise this exact demotion path. + mod = @eval module $(gensym()) + @noinline produce(cond::Bool, a::Int32) = cond ? a : Int64(7) + function kernel(ptr::Core.LLVMPtr{Int64,1}, cond::Bool, a::Int32) + x = produce(cond, a) + Base.unsafe_store!(ptr, x isa Int64 ? x : Int64(x)) + return + end + end + if LLVM.version() >= v"17" + # the box must be demoted, leaving a stack alloca rather than an addrspace(2) constant + air = sprint() do io + Metal.code_native(io, mod.kernel, (Core.LLVMPtr{Int64,1}, Bool, Int32); kernel=true) + end + @test occursin("alloca", air) + @test !occursin("addrspace(2) constant", air) + elseif GPUCompiler.supports_relocatable_ir() + # demotion needs LLVM.jl's `convert_users_to_instructions!` (LLVM 17+); the kernel + # must be rejected rather than silently miscompiled + @test_throws "require Julia 1.12" begin + sprint() do io + Metal.code_native(io, mod.kernel, (Core.LLVMPtr{Int64,1}, Bool, Int32); + kernel=true) + end + end + end + # on 1.10 codegen embeds the box as a raw host address instead of a module global, + # so there is nothing to demote (or to reject) +end + # Tuples with a dynamic index are lowered to an addrspace(2) constant plus a # GEP+load. Without InferAddressSpaces propagating AS 2 through the cast to # the generic AS introduced during `add_global_address_spaces!`, the load diff --git a/test/native.jl b/test/native.jl index 822c06e0..76358481 100644 --- a/test/native.jl +++ b/test/native.jl @@ -59,7 +59,9 @@ end end end - job, _ = Native.create_job(mod.outer, (Int, Symbol); validate=false) + # A relocatable back-end keeps the Symbol reference symbolic in `:llvm`. + job, _ = Native.create_job(mod.outer, (Int, Symbol); validate=false, + relocations=:patch) JuliaContext() do ctx ir, meta = GPUCompiler.compile(:llvm, job) @@ -72,18 +74,10 @@ end @test length(other_mis) == 1 @test only(other_mis).def in methods(mod.inner) - if VERSION >= v"1.12" - @test length(meta.gv_to_value) == 1 - for (k, v) in meta.gv_to_value - @test v != C_NULL - end + if GPUCompiler.supports_relocatable_ir() + @test length(meta.relocations) == 1 + @test only(meta.relocations.records).target isa GPUCompiler.JuliaValueRef end - # TODO: Global values get privatized, so we can't find them by name anymore. - # %.not = icmp eq ptr %"sym::Symbol", inttoptr (i64 140096668482288 to ptr), !dbg !38 - # for (name, v) in meta.gv_to_value - # gv = globals(ir)[name] - # @test LLVM.initializer(gv) === v - # end end end @@ -225,21 +219,6 @@ end @test new_res !== res @test new_res.asm === nothing - @static if GPUCompiler.HAS_INTEGRATED_CACHE - # session-dependent results (e.g. artifacts with relocated GVs) are wiped - # before image serialization; emulate the atexit-driven wipe directly - new_res.asm = "session-dependent" - other_job, _ = Native.create_job(mod.kernel, (Int64,); name="other") - other_res = GPUCompiler.cached_results(mod.Results, other_job) - push!(GPUCompiler.session_dependent_jobs, new_job) - GPUCompiler.wipe_session_dependent_results() - @test isempty(GPUCompiler.session_dependent_jobs) - wiped_res = GPUCompiler.cached_results(mod.Results, new_job) - @test wiped_res !== new_res - @test wiped_res.asm === nothing - # ... without affecting other configs on the same CI - @test GPUCompiler.cached_results(mod.Results, other_job) === other_res - end end @testset "runtime cache invalidation" begin @@ -273,11 +252,10 @@ end end end - @testset "runtime constgv relocation" begin + @testset "runtime relocations" begin # runtime functions like `box_bool` may reference Julia singletons through - # `julia.constgv` globals. Their session-absolute addresses must be baked into - # the cached runtime bitcode when it is built: only kernel modules go through - # `relocate_gvs!`, so a slot left null here would stay null on the device. + # `julia.constgv` globals. Keep their Julia identities with the cached bitcode + # so the final kernel can resolve them in its own session. job, _ = Native.create_job(identity, (Nothing,)) JuliaContext() do ctx GPUCompiler.load_runtime(job) @@ -292,10 +270,12 @@ end isempty(uses(gv)) && continue used += 1 init = LLVM.initializer(gv) - @test init !== nothing && !LLVM.isnull(init) + @test init === nothing + rec = GPUCompiler.find_relocation(lib.relocations, LLVM.name(gv)) + @test rec !== nothing && rec.kind === GPUCompiler.SlotSite end - @static if VERSION >= v"1.12-" - # on older versions, Julia bakes addresses without tagging globals + if GPUCompiler.supports_relocatable_ir() + # otherwise Julia embeds addresses without tagging globals @test used > 0 end end @@ -369,13 +349,17 @@ end Native.code_execution(mod.kernel, (Ptr{Int64}, Bool, Int32)) Native.code_execution(mod.egal_kernel, (Ptr{Bool}, Bool, Int32)) - # relocate_gvs! reports whether the module stayed session-portable + # Classification records whether the module stayed session-portable; eager + # lowering then resolves any remaining relocation slots. + collect_job, _ = Native.create_job(mod.kernel, (Ptr{Int64}, Bool, Int32)) + namespace = GPUCompiler.relocation_namespace(collect_job) + collect!(m, map) = GPUCompiler.collect_julia_value_relocations!(collect_job, m, map) JuliaContext() do ctx # Unlike Int128, vector-shaped tuples are 16-byte aligned on all # supported architectures and Julia versions. aligned = (VecElement(Int64(1)), VecElement(Int64(2))) @test Base.datatype_alignment(typeof(aligned)) > sizeof(Int) - objs = Any[Int64(42), 1.25, :sym, aligned] + objs = Any[Int64(42), 1.25, :sym, aligned, Union{}] # pointers to the heap boxes rooted in `objs` (passing an element # through a specialized function would re-box, possibly on the stack) ptrs = [ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), x) for x in objs] @@ -392,9 +376,12 @@ end gv = LLVM.GlobalVariable(m, LLVM.PointerType(LLVM.Int8Type()), name) constant!(gv, true) end - @test GPUCompiler.relocate_gvs!(m, Dict{String, Ptr{Cvoid}}()) + relocs = collect!(m, Dict{String, Ptr{Cvoid}}()) + @test isempty(relocs) + GPUCompiler.bake_relocations!(m, relocs) bool_ir = string(m) for name in ("jl_true", "jl_false") + # a fully-materialized box is private, so it needs no per-job namespace @test haskey(globals(m), "$(name)_box") @test occursin("@$name = private constant", bool_ir) end @@ -405,29 +392,81 @@ end GC.@preserve objs begin # smalltag isbits: materialized, portable m, map = slot_module(ptrs[1]) - @test GPUCompiler.relocate_gvs!(m, map) + relocs = collect!(m, map) + @test isempty(relocs) + GPUCompiler.bake_relocations!(m, relocs) @test haskey(globals(m), "jl_global_0_box") dispose(m) - # Float64: materialized, but the header carries a type pointer + # Float64: the non-smalltag header is an interior relocation. m, map = slot_module(ptrs[2]) - @test !GPUCompiler.relocate_gvs!(m, map) - @test haskey(globals(m), "jl_global_0_box") + relocs = collect!(m, map) + @test length(relocs) == 1 + rec = only(relocs.records) + @test rec.kind === GPUCompiler.InteriorSite + @test rec.offset == 0 + @test rec.target.value === Float64 + # a relocatable box is addressed by name, so its name carries the namespace + @test startswith(rec.name, namespace) + box = globals(m)[rec.name] + @test isextinit(box) + @test linkage(box) == LLVM.API.LLVMExternalLinkage + header_idx = Int(element_at(datalayout(m), global_value_type(box), + rec.offset)) + 1 + @test convert(UInt, collect(operands(initializer(box)))[header_idx]) == 0 + GPUCompiler.bake_relocations!(m, relocs) + @test isempty(relocs) + @test !isextinit(box) + @test isconstant(box) + @test linkage(box) == LLVM.API.LLVMPrivateLinkage + @test convert(UInt, collect(operands(initializer(box)))[header_idx]) == + GPUCompiler.resolve_relocation_target(rec.target) dispose(m) - # Symbol: baked address + # Symbol: resolved address m, map = slot_module(ptrs[3]) - @test !GPUCompiler.relocate_gvs!(m, map) + relocs = collect!(m, map) + rec = only(relocs.records) + @test rec.kind === GPUCompiler.SlotSite + @test rec.target.value === objs[3] + @test startswith(rec.name, namespace) + GPUCompiler.bake_relocations!(m, relocs) + @test isempty(relocs) @test !haskey(globals(m), "jl_global_0_box") @test occursin("inttoptr", string(m)) dispose(m) + # Empty type objects have a zero-sized singleton representation. + m, map = slot_module(ptrs[5]) + relocs = collect!(m, map) + rec = only(relocs.records) + @test rec.kind === GPUCompiler.SlotSite + @test rec.target.value === Union{} + dispose(m) + # 16-byte-aligned payloads get padded past the header word m, map = slot_module(ptrs[4]) - GPUCompiler.relocate_gvs!(m, map) - box = globals(m)["jl_global_0_box"] + relocs = collect!(m, map) + rec = only(relocs.records) + @test rec.offset == 8 + GPUCompiler.bake_relocations!(m, relocs) + box = globals(m)[rec.name] @test length(elements(LLVM.global_value_type(box))) == 3 dispose(m) + + # Codegen can emit several slots for one value in a module (observed on + # 1.11, whose backported GV API does not deduplicate); their + # content-derived names collide, so later slots must alias the first. + m = LLVM.Module("duplicate slots") + gvs = [LLVM.GlobalVariable(m, LLVM.PointerType(LLVM.Int8Type()), + "jl_global#$i") for i in 1:2] + relocs = collect!(m, Dict("jl_global#1" => ptrs[3], + "jl_global#2" => ptrs[3])) + rec = only(relocs.records) + @test rec.target.value === objs[3] + @test count(gv -> startswith(LLVM.name(gv), namespace), globals(m)) == 1 + @test !any(gv -> startswith(LLVM.name(gv), "jl_global#"), globals(m)) + dispose(m) end end end @@ -733,12 +772,416 @@ end end end -@testset "CPU reference resolution" begin +@testset "relocation target resolution" begin + ref = GPUCompiler.JuliaValueRef(:probe) + @test_throws ArgumentError GPUCompiler.Relocation(GPUCompiler.SlotSite, "invalid", -1, ref) + # a slot is a whole word, so only an interior record may carry an offset + @test_throws ArgumentError GPUCompiler.Relocation(GPUCompiler.SlotSite, "slot", 8, ref) + + sym = :relocation_target_probe + @test GPUCompiler.resolve_relocation_target(GPUCompiler.JuliaValueRef(sym)) == + UInt(pointer_from_objref(sym)) + + singleton = nothing + @test GPUCompiler.resolve_relocation_target(GPUCompiler.JuliaValueRef(singleton)) == + UInt(ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), singleton)) + + @test_throws ErrorException GPUCompiler.JuliaValueRef(1.5) +end + +@testset "applied relocation execution" begin + # A consumer that resolves the metadata into a module of its own instead of letting a + # loader do it: it caches the `:llvm` result plus the relocation metadata, and in every + # later session re-parses, `apply_relocations!`s, and emits an object with nothing left + # symbolic. AllocCheck does exactly this for the module it analyzes. + if GPUCompiler.supports_relocatable_ir() + mod = @eval module $(gensym()) + # the boxed `1.0` alternative of the isbits union is an *interior* record (its + # header word is a `Float64` type tag), while the Symbol is a whole-word slot; + # `apply_relocations!` must handle both kinds + @noinline produce(cond::Bool, a::Int32) = cond ? a : 1.0 + function f(cond::Bool) + x = produce(cond, Int32(7)) + word = x isa Float64 ? reinterpret(UInt64, x) : UInt64(0) + return word + UInt(pointer_from_objref(:applied_probe)) + end + end + job, _ = Native.create_job(mod.f, (Bool,); relocations=:patch) + JuliaContext() do ctx + ir, meta = GPUCompiler.compile(:llvm, job) + relocs = meta.relocations + @test all(rec -> rec.target isa GPUCompiler.JuliaValueRef, relocs.records) + # both record kinds are exercised + @test any(rec -> rec.kind === GPUCompiler.SlotSite, relocs.records) + @test any(rec -> rec.kind === GPUCompiler.InteriorSite, relocs.records) + + # the cache artifact: session-portable bitcode + relocation metadata + bitcode = let io = IOBuffer() + write(io, ir) + take!(io) + end + entry = LLVM.name(meta.entry) + + # a fresh session resolves the records into its own copy of the module, and only + # then emits an object + session_mod = parse(LLVM.Module, MemoryBuffer(bitcode)) + GPUCompiler.apply_relocations!(session_mod, relocs) + @test !isempty(relocs) # the metadata is not consumed + obj, _ = GPUCompiler.emit_asm(job, session_mod, LLVM.API.LLVMObjectFile) + + expected = reinterpret(UInt64, 1.0) + + GPUCompiler.resolve_relocation_target( + GPUCompiler.JuliaValueRef(:applied_probe)) + fptr, lljit, _table = Native.load(Vector{UInt8}(codeunits(obj)), entry, + GPUCompiler.Relocations()) + try + # the boxed alternative: its header tag decides the `isa`, so a stranded + # interior record would show up as a wrong result rather than a crash + @test ccall(fptr, UInt, (Bool,), false) == expected + @test ccall(fptr, UInt, (Bool,), false) == mod.f(false) + # ...and the inline alternative, which only reads the Symbol slot + @test ccall(fptr, UInt, (Bool,), true) == mod.f(true) + finally + dispose(lljit) + end + end + end +end + +@testset "eager relocation resolution" begin + # Eager resolution in `emit_llvm` leaves nothing for a loader. + mod = @eval module $(gensym()) + probe() = UInt(pointer_from_objref(:eager_probe)) + end + job, _ = Native.create_job(mod.probe, Tuple{}) + JuliaContext() do ctx + ir, meta = GPUCompiler.compile(:llvm, job) + @test isempty(meta.relocations) + # nothing is left for a loader to patch or import + @test !any(GPUCompiler.isextinit, globals(ir)) + + # This back-end can emit objects without threading relocation metadata. + code, _ = GPUCompiler.emit_asm(job, ir, LLVM.API.LLVMObjectFile) + @test !isempty(code) + end +end + +@testset "patchable relocation" begin + # An object-caching consumer: the words are written into the loaded image, so a cached + # object needs no compiler at all in a later session. CUDA does this with + # `cuModuleGetGlobal` + `cuMemcpyHtoD`; the ORC JIT below proves the same works under + # JITLink on macOS aarch64, the strictest W^X environment (only code pages are hardened). + if GPUCompiler.supports_relocatable_ir() + mod = @eval module $(gensym()) + f() = UInt(pointer_from_objref(:patch_probe)) + end + job, _ = Native.create_job(mod.f, Tuple{}; relocations=:patch) + JuliaContext() do ctx + obj, meta = GPUCompiler.compile(:obj, job) + relocs = meta.relocations + @test !isempty(relocs) + + # every slot became a null-init, externally-initialized definition kept alive by + # `llvm.used`; the loader patches each record after loading. Definitions are weak + # so that two objects defining one record coalesce (see "shared patchable record"). + @test haskey(globals(meta.ir), "llvm.used") + for rec in relocs.records + gv = globals(meta.ir)[rec.name] + @test !isdeclaration(gv) + @test isextinit(gv) + @test !isconstant(gv) + @test linkage(gv) == LLVM.API.LLVMWeakODRLinkage + rec.kind === GPUCompiler.SlotSite && @test LLVM.isnull(initializer(gv)) + end + + bytes = Vector{UInt8}(codeunits(obj)) + entry = LLVM.name(meta.entry) + probe = only(filter(rec -> rec.target isa GPUCompiler.JuliaValueRef && + rec.target.value === :patch_probe, relocs.records)) + expected = GPUCompiler.resolve_relocation_target(probe.target) + fptr, lljit, _table = Native.load(bytes, entry, relocs) + try + @test ccall(fptr, UInt, ()) == expected + finally + dispose(lljit) + end + + # The manifest now describes an emitted object, so dropping a record would leave + # its definition holding a zero and mis-branch silently. Refuse instead. + @test_throws "already been lowered" GPUCompiler.prune_dead_relocations!( + meta.ir, relocs) + @test_throws "already been lowered" GPUCompiler.add_relocation!( + relocs, GPUCompiler.SlotSite, "late", 0, probe.target) + + # A consumer that also wants a session-resolved copy to analyze (AllocCheck reads + # type tags out of one) applies the manifest *after* `emit_asm` froze it. That + # works because resolution goes into a copy — which freezing could easily break. + GPUCompiler.apply_relocations!(meta.ir, relocs) + @test !isempty(relocs) + end + end +end + +@testset "shared patchable record" begin + # One record can be defined by two objects at once: a relocation-carrying runtime-library + # function keeps its own job's namespace in every kernel it is linked into. An + # AllocCheck-shaped loader puts every object in one JITDylib, where two definitions of a + # symbol is an error unless they are weak. Construct the collision directly so this does + # not depend on which runtime functions the test kernel happens to import. + if GPUCompiler.supports_relocatable_ir() + mod = @eval module $(gensym()) + f() = 0 + end + job, _ = Native.create_job(mod.f, Tuple{}; relocations=:patch) + JuliaContext() do ctx + ptr(T) = GPUCompiler.supports_typed_pointers(ctx) ? "$T*" : "ptr" + ref = GPUCompiler.JuliaValueRef(:shared_probe) + function shared_object(entry) + m = parse(LLVM.Module, """ + @shared_reloc = external global i64 + + define i64 @$entry() { + %value = load i64, $(ptr("i64")) @shared_reloc + ret i64 %value + }""") + relocs = GPUCompiler.Relocations( + [GPUCompiler.Relocation(GPUCompiler.SlotSite, "shared_reloc", 0, ref)]) + asm, _ = GPUCompiler.emit_asm(job, m, relocs, LLVM.API.LLVMObjectFile) + @test linkage(globals(m)["shared_reloc"]) == LLVM.API.LLVMWeakODRLinkage + return Vector{UInt8}(codeunits(asm)), relocs + end + + obj_a, relocs_a = shared_object("shared_entry_a") + obj_b, _ = shared_object("shared_entry_b") + expected = GPUCompiler.resolve_relocation_target(ref) + + lljit = LLJIT(; tm=JITTargetMachine()) + try + jd = JITDylib(lljit) + add!(lljit, jd, MemoryBuffer(obj_a)) + add!(lljit, jd, MemoryBuffer(obj_b)) # the duplicate definition + + # patching the one surviving definition serves both objects + for (rec, word) in GPUCompiler.resolved_relocations(relocs_a) + addr = lookup(lljit, rec.name) + unsafe_store!(Ptr{UInt}(pointer(addr) + rec.offset), word) + end + for entry in ("shared_entry_a", "shared_entry_b") + @test ccall(pointer(lookup(lljit, entry)), UInt, ()) == expected + end + finally + dispose(lljit) + end + end + end +end + +@testset "tabulated relocation" begin + # A consumer with no access at all to loaded code: every record is rewritten into an + # indexed load from a table of words the loader delivers as run-time data. Metal does + # this through the kernel state; the test back-end reaches the table through a single + # patchable global, so the strategy is covered off-device. + if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17" + mod = @eval module $(gensym()) + # both record kinds: an interior box header (the `Float64` tag) and a slot + @noinline produce(cond::Bool, a::Int32) = cond ? a : 2.0 + function f(cond::Bool) + x = produce(cond, Int32(7)) + word = x isa Float64 ? reinterpret(UInt64, x) : UInt64(0) + return word + UInt(pointer_from_objref(:table_probe)) + end + end + job, _ = Native.create_job(mod.f, (Bool,); relocations=:table) + JuliaContext() do ctx + obj, meta = GPUCompiler.compile(:obj, job) + relocs = meta.relocations + @test !isempty(relocs) + @test any(rec -> rec.kind === GPUCompiler.SlotSite, relocs.records) + @test any(rec -> rec.kind === GPUCompiler.InteriorSite, relocs.records) + + # every record's global is gone: slots are erased, boxes demoted to allocas + for rec in relocs.records + @test !haskey(globals(meta.ir), rec.name) + end + # the words are read out of the table, not baked into the module + @test haskey(globals(meta.ir), Native.RELOC_TABLE_BASE) + + expected = reinterpret(UInt64, 2.0) + + GPUCompiler.resolve_relocation_target( + GPUCompiler.JuliaValueRef(:table_probe)) + fptr, lljit, table = Native.load(Vector{UInt8}(codeunits(obj)), + LLVM.name(meta.entry), relocs; table=true) + try + GC.@preserve table begin + @test ccall(fptr, UInt, (Bool,), false) == expected + @test ccall(fptr, UInt, (Bool,), false) == mod.f(false) + @test ccall(fptr, UInt, (Bool,), true) == mod.f(true) + end + finally + dispose(lljit) + end + + # A record's index is its rank in the manifest, and that index is baked into the + # emitted code — so recompiling the same kernel must produce the same manifest in + # the same order, or a cached object and a freshly-resolved table would disagree. + # (This is what makes a relocation-carrying kernel's *bytes* stable, which the + # Metal.jl suite asserts on a real cache key.) + _, again = GPUCompiler.compile(:obj, job) + @test [(rec.kind, rec.name, rec.offset) for rec in again.relocations.records] == + [(rec.kind, rec.name, rec.offset) for rec in relocs.records] + + # The delivered words follow the order the lowering fixed, not the record vector, + # so losing records can no longer renumber the table. Pruning a copy drops *every* + # record here (the lowering erased all their globals) and the words still stand. + mutated = copy(relocs) + GPUCompiler.prune_dead_relocations!(meta.ir, mutated) + @test isempty(mutated) + @test GPUCompiler.resolved_relocation_table(mutated) == + GPUCompiler.resolved_relocation_table(relocs) + + # And the manifest itself refuses to be renumbered at all. + @test_throws "already been lowered" GPUCompiler.prune_dead_relocations!( + meta.ir, relocs) + end + end +end + +@testset "unlowered relocation table" begin + # Emitting a `:table` module through the 3-argument `emit_asm` hands the lowering an + # empty manifest, leaving the real one unlowered and the module's slots stranded. The + # loader is the first thing to notice, so it must say so rather than deliver no words. + if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17" + mod = @eval module $(gensym()) + f() = UInt(pointer_from_objref(:unlowered_probe)) + end + job, _ = Native.create_job(mod.f, Tuple{}; relocations=:table) + JuliaContext() do ctx + ir, meta = GPUCompiler.compile(:llvm, job) + @test !isempty(meta.relocations) + GPUCompiler.emit_asm(job, ir, LLVM.API.LLVMObjectFile) # the 3-arg form + @test_throws "never rewritten" GPUCompiler.resolved_relocation_table( + meta.relocations) + end + end +end + +@testset "relocation-free tabulated module" begin + # A module without relocations must not gain any table access at all: the strategy is + # free for the (overwhelmingly common) relocation-free kernel. + if LLVM.version() >= v"17" + mod = @eval module $(gensym()) + f(x::Int) = x + 1 + end + job, _ = Native.create_job(mod.f, (Int,); relocations=:table) + JuliaContext() do ctx + _, meta = GPUCompiler.compile(:obj, job) + @test isempty(meta.relocations) + @test !haskey(globals(meta.ir), Native.RELOC_TABLE_BASE) + end + end +end + +@testset "relocation validation errors" begin + JuliaContext() do ctx + word() = GPUCompiler.relocation_word_type() + nop = (_rec, _gv) -> nothing + ref = GPUCompiler.JuliaValueRef(:probe) + reloc(kind, name, offset=0) = GPUCompiler.Relocations( + [GPUCompiler.Relocation(kind, name, offset, ref)]) + slot(name) = reloc(GPUCompiler.SlotSite, name) + interior(name, offset) = reloc(GPUCompiler.InteriorSite, name, offset) + + # a record whose global is absent from the module + mod = LLVM.Module("errors") + @test_throws "Missing relocation global" GPUCompiler.foreach_relocation( + nop, mod, slot("absent")) + + # a slot must be word-sized + mod = LLVM.Module("errors") + GlobalVariable(mod, LLVM.Int32Type(), "narrow") + @test_throws "has size" GPUCompiler.foreach_relocation(nop, mod, slot("narrow")) + + # an interior record must name a definition... + mod = LLVM.Module("errors") + GlobalVariable(mod, word(), "decl") + @test_throws "is a declaration" GPUCompiler.foreach_relocation( + nop, mod, interior("decl", 0)) + + # ...whose initializer is a struct... + mod = LLVM.Module("errors") + gv = GlobalVariable(mod, word(), "flat") + initializer!(gv, ConstantInt(word(), 0)) + @test_throws "non-struct initializer" GPUCompiler.foreach_relocation( + nop, mod, interior("flat", 0)) + + # ...and it must land within that global + mod = LLVM.Module("errors") + gv = GlobalVariable(mod, LLVM.StructType([LLVM.Int64Type(), LLVM.Int64Type()]), "box") + initializer!(gv, ConstantStruct(LLVM.Constant[ConstantInt(0), ConstantInt(0)])) + @test_throws "outside its" GPUCompiler.foreach_relocation( + nop, mod, interior("box", 16)) + end +end + +@testset "prune dead relocations" begin + JuliaContext() do ctx + ptr(T) = GPUCompiler.supports_typed_pointers(ctx) ? "$T*" : "ptr" + mod = parse(LLVM.Module, """ + @live = external global i64 + @dead = internal global { i64, i64 } { i64 0, i64 0 } + define i64 @use() { + %v = load i64, $(ptr("i64")) @live + ret i64 %v + }""") + relocs = GPUCompiler.Relocations() + for (name, kind) in ("live" => GPUCompiler.SlotSite, + "dead" => GPUCompiler.InteriorSite, # unused definition + "absent" => GPUCompiler.SlotSite) # global already gone + GPUCompiler.add_relocation!(relocs, kind, name, 0, + GPUCompiler.JuliaValueRef(Symbol(name))) + end + GPUCompiler.prune_dead_relocations!(mod, relocs) + @test [rec.name for rec in relocs.records] == ["live"] + @test haskey(globals(mod), "live") # a used declaration survives + @test !haskey(globals(mod), "dead") # the dead definition is erased + end +end + +@testset "resolve zeroinitializer box" begin + # An all-zero box (a patchable header over a zero payload) is folded by LLVM to a + # `zeroinitializer`, a ConstantAggregateZero that reports no operands; resolution must + # resolve its header word. Regresses JuliaGPU/oneAPI.jl's "#55: invalid integers created + # by alloc_opt", where `SVector(0f0, 0f0)` boxed a zero payload. + JuliaContext() do ctx + mod = parse(LLVM.Module, + "@zero_box = private global { i64, [8 x i8] } zeroinitializer") + gv = globals(mod)["zero_box"] + @test initializer(gv) isa LLVM.ConstantAggregateZero # the folded shape + relocs = GPUCompiler.Relocations( + [GPUCompiler.Relocation(GPUCompiler.InteriorSite, "zero_box", 0, + GPUCompiler.JuliaValueRef(Float64))]) + GPUCompiler.bake_relocations!(mod, relocs) + init = initializer(gv) + @test !(init isa LLVM.ConstantAggregateZero) # rebuilt into explicit fields + header = convert(UInt, LLVM.Constant[operands(init)...][1]) + @test header == GPUCompiler.resolve_relocation_target(GPUCompiler.JuliaValueRef(Float64)) + @test isconstant(gv) + @test isempty(relocs) + end +end + +@testset "cglobal relocation" begin # JIT-private symbols like `jl_get_pgcstack_resolved` (JuliaLang/julia#61527) cannot # be looked up using `jl_cglobal`, so we should only resolve bindings that are # actually loaded from, leaving called functions alone. job, _ = Native.create_job(identity, (Nothing,)) JuliaContext() do ctx + ptr(T) = GPUCompiler.supports_typed_pointers(ctx) ? "$T*" : "ptr" + word_ptr = ptr("i8") + word_ptr_ptr = ptr(word_ptr) + function_word_ptr(name) = GPUCompiler.supports_typed_pointers(ctx) ? + "i64* bitcast (i64 ()* @$name to i64*)" : "ptr @$name" + mod = parse(LLVM.Module, """ declare void @jl_get_pgcstack_resolved() @@ -748,6 +1191,155 @@ end }""") GPUCompiler.prepare_execution!(job, mod) @test haskey(functions(mod), "jl_get_pgcstack_resolved") + + mod = parse(LLVM.Module, """ + @jl_float32_type = external global $word_ptr + + define $word_ptr @entry() { + %value = load $word_ptr, $word_ptr_ptr @jl_float32_type + ret $word_ptr %value + }""") + GPUCompiler.prepare_execution!(job, mod) + ir = string(mod) + @test !occursin("load $word_ptr, $word_ptr_ptr @jl_float32_type", ir) + expected = GPUCompiler.resolve_relocation_target( + GPUCompiler.CGlobalRef(:jl_float32_type)) + @test occursin("inttoptr (i64 $expected to $word_ptr)", ir) + + mod = parse(LLVM.Module, """ + @jl_float32_type = external global $word_ptr + + define $word_ptr @entry() { + %value = load $word_ptr, $word_ptr_ptr @jl_float32_type + ret $word_ptr %value + }""") + relocs = GPUCompiler.Relocations() + @test GPUCompiler.collect_cglobal_relocations!(job, mod, relocs) + rec = only(relocs.records) + @test rec.target == GPUCompiler.CGlobalRef(:jl_float32_type) + @test rec.kind === GPUCompiler.SlotSite + @test rec.offset == 0 + # the slot is a symbol a loader addresses, so its name carries the job's namespace + @test startswith(rec.name, GPUCompiler.relocation_namespace(job)) + @test occursin("@$(rec.name) = external global i64", string(mod)) + GPUCompiler.emit_patchable_relocations!(mod, relocs) + @test occursin("externally_initialized global i64 0", string(mod)) + + mod = parse(LLVM.Module, """ + declare i64 @jl_float32_type() + + define i64 @entry() { + %value = load i64, $(function_word_ptr("jl_float32_type")) + ret i64 %value + }""") + relocs = GPUCompiler.Relocations() + @test GPUCompiler.collect_cglobal_relocations!(job, mod, relocs) + rec = only(relocs.records) + @test rec.target == GPUCompiler.CGlobalRef(:jl_float32_type) + @test occursin("@$(rec.name) = external global i64", string(mod)) + GPUCompiler.emit_patchable_relocations!(mod, relocs) + @test occursin("externally_initialized global i64 0", string(mod)) + end +end + +@testset "relocation linking" begin + JuliaContext() do ctx + ptr(T) = GPUCompiler.supports_typed_pointers(ctx) ? "$T*" : "ptr" + + function slot_module(name, entry) + parse(LLVM.Module, """ + @$name = external global i64 + + define i64 @$entry() { + %value = load i64, $(ptr("i64")) @$name + ret i64 %value + }""") + end + + slot_relocs(name, value) = GPUCompiler.Relocations( + [GPUCompiler.Relocation(GPUCompiler.SlotSite, name, 0, + GPUCompiler.JuliaValueRef(value))]) + + # Per-job namespacing means a name never denotes two different words. Metadata for the + # same word still merges, though — the runtime library is linked into every kernel. + dest = slot_module("slot", "first") + dest_relocs = slot_relocs("slot", :shared) + src = slot_module("slot", "second") + src_relocs = slot_relocs("slot", :shared) + GPUCompiler.link_relocatable!(dest, dest_relocs, src, src_relocs) + @test [rec.name for rec in dest_relocs.records] == ["slot"] + @test occursin("@slot = external global i64", string(dest)) + + # Conflicting metadata for one word is an inconsistency, not a merge. + dest = slot_module("slot", "first") + dest_relocs = slot_relocs("slot", :first) + src = slot_module("slot", "second") + src_relocs = slot_relocs("slot", :second) + @test_throws "conflicting values" GPUCompiler.link_relocatable!( + dest, dest_relocs, src, src_relocs) + + # ...as is disagreement about what the word even is. + dest_relocs = slot_relocs("slot", :shared) + src_relocs = GPUCompiler.Relocations( + [GPUCompiler.Relocation(GPUCompiler.InteriorSite, "slot", 0, + GPUCompiler.JuliaValueRef(:shared))]) + @test_throws "recorded as both" GPUCompiler.add_relocation!( + dest_relocs, only(src_relocs.records)) + + # `only_needed` must keep metadata for imported slots and discard metadata for + # source globals that the LLVM linker did not import. + dest = parse(LLVM.Module, """ + declare i64 @source() + + define i64 @entry() { + %value = call i64 @source() + ret i64 %value + }""") + src = parse(LLVM.Module, """ + @used = external global i64 + @unused = external global i64 + + define i64 @source() { + %value = load i64, $(ptr("i64")) @used + ret i64 %value + }""") + src_relocs = GPUCompiler.Relocations() + for name in ("used", "unused") + GPUCompiler.add_relocation!(src_relocs, GPUCompiler.SlotSite, name, 0, + GPUCompiler.JuliaValueRef(Symbol(name))) + end + dest_relocs = GPUCompiler.Relocations() + GPUCompiler.link_relocatable!(dest, dest_relocs, src, src_relocs; + only_needed=true) + @test [rec.name for rec in dest_relocs.records] == ["used"] + @test only(dest_relocs.records).target.value === :used + + # Metadata for interior globals not imported under `only_needed` is discarded too. + dest = parse(LLVM.Module, """ + declare i64 @source_patch() + define i64 @entry_patch() { + %value = call i64 @source_patch() + ret i64 %value + }""") + src = parse(LLVM.Module, """ + @used_patch = externally_initialized global { i64, i64 } { i64 0, i64 1 } + + define i64 @source_patch() { + %value = load i64, $(ptr("i64")) getelementptr ({ i64, i64 }, $(ptr("{ i64, i64 }")) @used_patch, i32 0, i32 1) + ret i64 %value + }""") + unused = GlobalVariable(src, LLVM.StructType([LLVM.Int64Type(), LLVM.Int64Type()]), + "unused_patch") + initializer!(unused, ConstantStruct(LLVM.Constant[ConstantInt(0), ConstantInt(1)])) + src_relocs = GPUCompiler.Relocations() + for (name, T) in ("used_patch" => Float64, "unused_patch" => Int64) + GPUCompiler.add_relocation!(src_relocs, GPUCompiler.InteriorSite, name, 0, + GPUCompiler.JuliaValueRef(T)) + end + dest_relocs = GPUCompiler.Relocations() + GPUCompiler.link_relocatable!(dest, dest_relocs, src, src_relocs; + only_needed=true) + @test [rec.name for rec in dest_relocs.records] == ["used_patch"] end end @@ -1102,6 +1694,32 @@ end @test occursin("call void @julia_kernel", ir) end +@testset "Mock Enzyme deferred relocations" begin + # A deferred child that references a Julia value produces its own relocations; those + # must merge into the parent's metadata when the child module is linked in. + mod = @eval module $(gensym()) + import ..Enzyme + child(sym::Symbol) = sym === :deferred_reloc ? 1 : 2 + function parent(sym::Symbol) + ptr = Enzyme.deferred_codegen(typeof(child), Tuple{Symbol}) + return ccall(ptr, Int, (Symbol,), sym) + end + end + + # Keep the merged relocation symbolic so we can inspect it. + job, _ = Native.create_job(mod.parent, (Symbol,); relocations=:patch, validate=false) + JuliaContext() do ctx + ir, meta = GPUCompiler.compile(:llvm, job) + @test !occursin("deferred_codegen", string(ir)) + if GPUCompiler.supports_relocatable_ir() + @test any(meta.relocations.records) do rec + rec.target isa GPUCompiler.JuliaValueRef && + rec.target.value === :deferred_reloc + end + end + end +end + @testset "stack allocation intrinsic" begin mod = @eval module $(gensym()) import ..GPUCompiler diff --git a/test/native/precompile.jl b/test/native/precompile.jl index eb9abc4f..042ed1cc 100644 --- a/test/native/precompile.jl +++ b/test/native/precompile.jl @@ -37,22 +37,21 @@ precompile_test_harness("Inference caching") do load_path Results() = new(nothing) end - portable_kernel(x) = x + 1 + persistent_kernel(x) = x + 1 session_kernel(x) = x + 2 - # Attach representative back-end artifacts while the package image is built. The - # portable entry should survive serialization; the session-dependent one should be - # removed by GPUCompiler's pre-output atexit hook. + # A back-end implementing relocation keeps results across sessions. let - job, _ = NativeCompiler.Native.create_job(portable_kernel, (Int,)) + job, _ = NativeCompiler.Native.create_job(persistent_kernel, (Int,); relocations=:patch) precompile(job) - NativeCompiler.GPUCompiler.cached_results(Results, job).artifact = "portable" + NativeCompiler.GPUCompiler.cached_results(Results, job).artifact = "persistent" end + + # Default back-ends use the session-local store. let job, _ = NativeCompiler.Native.create_job(session_kernel, (Int,)) precompile(job) NativeCompiler.GPUCompiler.cached_results(Results, job).artifact = "session" - NativeCompiler.GPUCompiler.mark_session_dependent!(job) end let @@ -111,10 +110,15 @@ precompile_test_harness("Inference caching") do load_path using NativeBackend - portable_job, _ = NativeCompiler.Native.create_job(NativeBackend.portable_kernel, (Int,)) - portable_res = GPUCompiler.cached_results(NativeBackend.Results, portable_job) - @test portable_res !== nothing - @test portable_res.artifact == "portable" + persistent_job, _ = NativeCompiler.Native.create_job( + NativeBackend.persistent_kernel, (Int,); relocations=:patch) + persistent_res = GPUCompiler.cached_results(NativeBackend.Results, persistent_job) + @test persistent_res !== nothing + if GPUCompiler.supports_relocatable_ir() + @test persistent_res.artifact == "persistent" + else + @test persistent_res.artifact === nothing + end session_job, _ = NativeCompiler.Native.create_job(NativeBackend.session_kernel, (Int,)) session_res = GPUCompiler.cached_results(NativeBackend.Results, session_job) diff --git a/test/ptx.jl b/test/ptx.jl index e98f55d8..3ab82d07 100644 --- a/test/ptx.jl +++ b/test/ptx.jl @@ -29,19 +29,35 @@ end end @testset "global variable relocation" begin - # references to Julia objects (`julia.constgv` globals, e.g. Symbol literals) must - # survive until `relocate_gvs!` bakes in their addresses at the toplevel link step. - # they used to be kept alive as internal globals with a null initializer, which the - # GlobalOpt run in `finish_module!` folded away, constant-folding any comparison - # against them (JuliaGPU/CUDA.jl#3185: kernels specialized on Symbols misbehaved). + # Julia-object references must remain declarations until relocation lowering. Null + # definitions would let GlobalOpt fold comparisons against them. mod = @eval module $(gensym()) kernel(name::Symbol) = name === :var ? 1 : 2 end + # `patch=true` keeps the slot symbolic instead of resolving it into the IR. ir = sprint() do io - PTX.code_llvm(io, mod.kernel, Tuple{Symbol}; dump_module=true) + PTX.code_llvm(io, mod.kernel, Tuple{Symbol}; dump_module=true, patch=true) end - addr = UInt64(pointer_from_objref(:var)) - @test occursin(string(addr), ir) + # Julia embeds the Symbol pointer before exposing the IR to GPUCompiler when it + # can't emit relocatable global metadata; otherwise it stays a symbolic slot we preserve. + # the slot keeps its content-derived name under the per-job namespace prefix + if GPUCompiler.supports_relocatable_ir() + @test occursin(r"jl_sym_var_[0-9a-f]+\"? = external", ir) + else + @test occursin(r"jl_sym_var_[0-9a-f]+\"? = external", ir) || occursin("inttoptr", ir) + end + @test !occursin("@\"jl_sym#", ir) +end + +@testset "Julia value global names" begin + # Julia's external codegen can name the declaration for this Symbol with `#`. + # The final PTX path must see the sanitized Julia value global, not the original + # declaration name rejected by NVPTX. + mod = @eval module $(gensym()) + const unusual_symbol = Symbol("value#global") + kernel(name::Symbol) = (name === unusual_symbol; return) + end + @test PTX.code_execution(mod.kernel, Tuple{Symbol}) !== nothing end @testset "boxed Bool singleton relocation" begin @@ -72,6 +88,28 @@ end end end +@testset "boxed header relocation" begin + mod = @eval module $(gensym()) + @noinline produce(cond::Bool, value::Int32) = cond ? value : 1.5 + function consume(cond::Bool, value::Int32) + x = produce(cond, value) + x isa Float64 && return x + return 0.0 + end + end + # `patch=true` keeps the interior header relocation symbolic (externally_initialized). + ir = sprint(io->PTX.code_llvm(io, mod.consume, Tuple{Bool,Int32}; + dump_module=true, patch=true)) + # As with Symbol literals above, Julia embeds the type pointer when it can't emit + # relocatable global metadata; otherwise GPUCompiler represents it as a relocation. + if GPUCompiler.supports_relocatable_ir() + @test occursin(r"@[A-Za-z0-9_]+_box = externally_initialized global", ir) + else + @test occursin(r"@[A-Za-z0-9_]+_box = externally_initialized global", ir) || + occursin("inttoptr", ir) + end +end + @testset "kernel functions" begin @testset "kernel argument attributes" begin mod = @eval module $(gensym()) @@ -200,6 +238,24 @@ if :NVPTX in LLVM.backends() end end +@testset "patchable relocation" begin + # A patching back-end (like CUDA.jl) keeps relocations symbolic; the patchable box must + # survive lowering into the generated PTX as a `.global` for the loader to write. + if GPUCompiler.supports_relocatable_ir() + mod = @eval module $(gensym()) + @noinline produce(cond::Bool, value::Int32) = cond ? value : 1.5 + function consume(cond::Bool, value::Int32) + x = produce(cond, value) + x isa Float64 && return x + return 0.0 + end + end + ptx = sprint(io->PTX.code_native(io, mod.consume, Tuple{Bool,Int32}; + dump_module=true, patch=true)) + @test occursin(r"\.global .*_box", ptx) + end +end + @testset "child functions" begin # we often test using @noinline child functions, so test whether these survive # (despite not having side-effects) diff --git a/test/spirv.jl b/test/spirv.jl index fac3ea2c..019a9280 100644 --- a/test/spirv.jl +++ b/test/spirv.jl @@ -45,6 +45,24 @@ end SPIRV.code_llvm(mod.kernel, Tuple{}; backend, kernel=true) end end + +@testset "baked boxed relocation cleanup" begin + mod = @eval module $(gensym()) + @noinline produce(cond::Bool, value::Int32) = cond ? value : 1.5 + function kernel(out::Core.LLVMPtr{UInt,1}, cond::Bool, value::Int32) + x = produce(cond, value) + Base.unsafe_store!(out, UInt(x isa Float64)) + return + end + end + + # Baking an interior relocation can expose a dead pointer component of an isbits-union + # result. It must be folded before SPIR-V translation, which otherwise emits a reference + # to the now-unused box without defining it. + _, meta = SPIRV.code_execution( + mod.kernel, (Core.LLVMPtr{UInt,1}, Bool, Int32); backend) + @test all(!endswith(LLVM.name(gv), "_box") for gv in globals(meta.ir)) +end end @testset "unsupported type detection" begin