Skip to content

Commit

Permalink
use object_build_id
Browse files Browse the repository at this point in the history
  • Loading branch information
vchuravy committed Apr 16, 2024
1 parent 397825d commit 13cae6b
Showing 1 changed file with 47 additions and 25 deletions.
72 changes: 47 additions & 25 deletions src/execution.jl
Original file line number Diff line number Diff line change
Expand Up @@ -78,14 +78,6 @@ end

cache_path() = @get_scratch!("cache")
clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
function cache_path(key)
return joinpath(
cache_path(),
# TODO: Use object_build_id from https://github.com/JuliaLang/julia/pull/53943
# Should we disk cache "runtime compilation".
string(Base.module_build_id(GPUCompiler)), # captures dependencies as well
string(cache_key), "ir.jls")
end

const cache_lock = ReentrantLock()

Expand Down Expand Up @@ -133,6 +125,30 @@ function cached_compilation(cache::AbstractDict{<:Any,V},
return obj::V
end

@noinline function cache_file(ci::CodeInstance, cfg::CompilerConfig)
@static if isdefined(Base, :object_build_id)
id = Base.object_build_id(ci)
if id === nothing # CI is from a runtime compilation, not worth caching on disk
return nothing
else
id = id % UInt64 # The upper 64bit are a checksum, unavailable during precompilation
end
else
id = Base.objectid(ci)
end

gpucompiler_buildid = Base.module_build_id(@__MODULE__)
if (gpucompiler_buildid >> 64) % UInt64 == 0xffffffffffffffff
return nothing # Don't cache during precompilation of GPUCompiler
end

return joinpath(
cache_path(),
# bifurcate the cache by build id of GPUCompiler
string(gpucompiler_buildid),
string(hash(cfg, hash(id)), ".jls"))
end

@noinline function actual_compilation(cache::AbstractDict, src::MethodInstance, world::UInt,
cfg::CompilerConfig, compiler::Function, linker::Function)
job = CompilerJob(src, cfg, world)
Expand All @@ -148,18 +164,22 @@ end
# slow path: compile and link
if obj === nothing || compile_hook[] !== nothing
asm = nothing
@static if VERSION >= v"1.11.0-" && disk_cache()
cache_key = Base.objectid(ci)
path = cache_path(cache_key)
if isfile(path)
try
@debug "Loading compiled kernel for $spec from $path"
asm = deserialize(path)
catch ex
@warn "Failed to load compiled kernel at $path" exception=(ex, catch_backtrace())
path = nothing
ondisk_hit = false
@static if VERSION >= v"1.11.0-"
if disk_cache() # TODO: (Should we allow backends to opt out?)
path = cache_file(ci, cfg)
if path !== nothing && isfile(path)
ondisk_hit = true
try
@debug "Loading compiled kernel for $spec from $path"
asm = deserialize(path)
catch ex
@warn "Failed to load compiled kernel at $path" exception=(ex, catch_backtrace())
end
end
end
else
end

if asm === nothing
asm = compiler(job)
Expand All @@ -169,13 +189,15 @@ end
return obj
end

@static if VERSION >= v"1.11.0-" && disk_cache() && !isfile(path)
# TODO: Should we only write out during precompilation?
tmppath, io = mktemp(;cleanup=false)
serialize(io, asm)
close(io)
# atomic move
Base.rename(tmppath, path, force=true)
@static if VERSION >= v"1.11.0-"
if !ondisk_hit && path !== nothing && disk_cache()
# TODO: Do we want to serialize some more metadata to make sure the asm matches?
tmppath, io = mktemp(;cleanup=false)
serialize(io, asm)
close(io)
# atomic move
Base.rename(tmppath, path, force=true)
end
end

obj = linker(job, asm)
Expand Down

0 comments on commit 13cae6b

Please sign in to comment.