Skip to content

Make generated code relocatable across Julia sessions #878

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
maleadt wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Make generated code relocatable across Julia sessions #878

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "GPUCompiler"
uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
version = "2.1.1"
version = "2.2.0"
authors = ["Tim Besard <tim.besard@gmail.com>"]

[workspace]
Expand Down Expand Up @@ -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"
Expand Down
8 changes: 3 additions & 5 deletions src/GPUCompiler.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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

Expand Down
40 changes: 1 addition & 39 deletions src/deprecated.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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


Expand Down
74 changes: 54 additions & 20 deletions src/driver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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]

Expand Down Expand Up @@ -405,29 +421,47 @@ 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

if should_verify()
@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
Loading