Skip to content

Commit d7196b6

Browse files
committed
Add disk cache infrastructure for Julia 1.11
1 parent c15433c commit d7196b6

File tree

5 files changed

+126
-15
lines changed

5 files changed

+126
-15
lines changed

Project.toml

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,23 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
99
LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
1010
Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
1111
Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
12+
Preferences = "21216c6a-2e73-6563-6e65-726566657250"
1213
Scratch = "6c6a2e73-6563-6170-7368-637461726353"
14+
Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
15+
TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
1316
TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
1417
UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
1518

1619
[compat]
1720
ExprTools = "0.1"
1821
InteractiveUtils = "1"
22+
LLVM = "6.6"
1923
Libdl = "1"
2024
Logging = "1"
21-
UUIDs = "1"
22-
LLVM = "6.6"
25+
Preferences = "1"
2326
Scratch = "1"
27+
Serialization = "1"
28+
TOML = "1"
2429
TimerOutputs = "0.5"
30+
UUIDs = "1"
2531
julia = "1.8"

src/GPUCompiler.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ using ExprTools: splitdef, combinedef
99

1010
using Libdl
1111

12+
using Serialization
1213
using Scratch: @get_scratch!
14+
using Preferences
1315

1416
const CC = Core.Compiler
1517
using Core: MethodInstance, CodeInstance, CodeInfo

src/execution.jl

Lines changed: 79 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,23 @@ end
6161

6262

6363
## cached compilation
64+
disk_cache() = parse(Bool, @load_preference("disk_cache", "false"))
65+
66+
"""
67+
enable_cache!(state::Bool=true)
68+
69+
Activate the GPUCompiler disk cache in the current environment.
70+
You will need to restart your Julia environment for it to take effect.
71+
72+
!!! note
73+
The cache functionality requires Julia 1.11
74+
"""
75+
function enable_cache!(state::Bool=true)
76+
@set_preferences!("disk_cache"=>string(state))
77+
end
78+
79+
cache_path() = @get_scratch!("cache")
80+
clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
6481

6582
const cache_lock = ReentrantLock()
6683

@@ -108,6 +125,30 @@ function cached_compilation(cache::AbstractDict{<:Any,V},
108125
return obj::V
109126
end
110127

128+
@noinline function cache_file(ci::CodeInstance, cfg::CompilerConfig)
129+
@static if isdefined(Base, :object_build_id)
130+
id = Base.object_build_id(ci)
131+
if id === nothing # CI is from a runtime compilation, not worth caching on disk
132+
return nothing
133+
else
134+
id = id % UInt64 # The upper 64bit are a checksum, unavailable during precompilation
135+
end
136+
else
137+
id = Base.objectid(ci)
138+
end
139+
140+
gpucompiler_buildid = Base.module_build_id(@__MODULE__)
141+
if (gpucompiler_buildid >> 64) % UInt64 == 0xffffffffffffffff
142+
return nothing # Don't cache during precompilation of GPUCompiler
143+
end
144+
145+
return joinpath(
146+
cache_path(),
147+
# bifurcate the cache by build id of GPUCompiler
148+
string(gpucompiler_buildid),
149+
string(hash(cfg, hash(id)), ".jls"))
150+
end
151+
111152
@noinline function actual_compilation(cache::AbstractDict, src::MethodInstance, world::UInt,
112153
cfg::CompilerConfig, compiler::Function, linker::Function)
113154
job = CompilerJob(src, cfg, world)
@@ -117,20 +158,53 @@ end
117158
ci = ci_cache_lookup(ci_cache(job), src, world, world)::Union{Nothing,CodeInstance}
118159
if ci !== nothing
119160
key = (ci, cfg)
120-
if haskey(cache, key)
121-
obj = cache[key]
122-
end
161+
obj = get(cache, key, nothing)
123162
end
124163

125164
# slow path: compile and link
126165
if obj === nothing || compile_hook[] !== nothing
127-
# TODO: consider loading the assembly from an on-disk cache here
128-
asm = compiler(job)
166+
asm = nothing
167+
path = nothing
168+
ondisk_hit = false
169+
@static if VERSION >= v"1.11.0-"
170+
# Don't try to hit the disk cache if we are for a *compile* hook
171+
if ci !== nothing && obj === nothing && disk_cache() # TODO: (Should we allow backends to opt out?)
172+
path = cache_file(ci, cfg)
173+
@debug "Looking for on-disk cache" job path
174+
if path !== nothing && isfile(path)
175+
ondisk_hit = true
176+
try
177+
@debug "Loading compiled kernel" job path
178+
asm = deserialize(path)
179+
catch ex
180+
@warn "Failed to load compiled kernel" job path exception=(ex, catch_backtrace())
181+
end
182+
end
183+
end
184+
end
129185

186+
if asm === nothing || compile_hook[] !== nothing
187+
# Run the compiler in-case we need to hook it.
188+
asm = compiler(job)
189+
end
130190
if obj !== nothing
131191
# we got here because of a *compile* hook; don't bother linking
132192
return obj
133193
end
194+
195+
@static if VERSION >= v"1.11.0-"
196+
if !ondisk_hit && path !== nothing && disk_cache()
197+
@debug "Writing out on-disk cache" job path
198+
# TODO: Do we want to serialize some more metadata to make sure the asm matches?
199+
tmppath, io = mktemp(;cleanup=false)
200+
serialize(io, asm)
201+
close(io)
202+
# atomic move
203+
mkpath(dirname(path))
204+
Base.rename(tmppath, path, force=true)
205+
end
206+
end
207+
134208
obj = linker(job, asm)
135209

136210
if ci === nothing

test/native_tests.jl

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -549,12 +549,13 @@ precompile_test_harness("Inference caching") do load_path
549549
import GPUCompiler
550550
using PrecompileTools
551551

552-
function kernel()
552+
function kernel(A, x)
553+
A[1] = x
553554
return
554555
end
555556

556557
let
557-
job, _ = NativeCompiler.create_job(kernel, ())
558+
job, _ = NativeCompiler.create_job(kernel, (Vector{Int}, Int))
558559
GPUCompiler.code_typed(job)
559560
end
560561

@@ -578,20 +579,30 @@ precompile_test_harness("Inference caching") do load_path
578579
job, _ = NativeCompiler.create_job(identity, (Int,))
579580
GPUCompiler.ci_cache_token(job)
580581
end
581-
ci = isdefined(identity_mi, :cache) ? identity_mi.cache : nothing
582-
while ci !== nothing
583-
@test ci.owner !== token
584-
ci = isdefined(ci, :next) ? ci.next : nothing
585-
end
582+
@test !check_presence(identity_mi, token)
586583

587584
using InferenceCaching
588585

589586
# Check that kernel survived
590-
kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{})
587+
kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{Vector{Int}, Int})
591588
@test check_presence(kernel_mi, token)
592589

593590
# check that identity survived
594591
@test check_presence(identity_mi, token)
592+
593+
GPUCompiler.enable_cache!()
594+
@test GPUCompiler.disk_cache() == true
595+
596+
job, _ = NativeCompiler.create_job(InferenceCaching.kernel, (Vector{Int}, Int))
597+
@assert job.source == kernel_mi
598+
ci = GPUCompiler.ci_cache_lookup(GPUCompiler.ci_cache(job), job.source, job.world, job.world)
599+
@assert ci !== nothing
600+
@assert ci.inferred !== nothing
601+
path = GPUCompiler.cache_file(ci, job.config)
602+
@test path !== nothing
603+
@test !ispath(path)
604+
NativeCompiler.cached_execution(InferenceCaching.kernel, (Vector{Int}, Int))
605+
@test ispath(path)
595606
end
596607
end
597608

test/native_testsetup.jl

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,4 +71,22 @@ function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
7171
end
7272
end
7373

74+
const runtime_cache = Dict{Any, Any}()
75+
76+
function compiler(job)
77+
JuliaContext() do ctx
78+
GPUCompiler.compile(:asm, job, validate=false)
79+
end
80+
end
81+
82+
function linker(job, asm)
83+
asm
84+
end
85+
86+
# simulates cached codegen
87+
function cached_execution(@nospecialize(func), @nospecialize(types); kwargs...)
88+
job, kwargs = create_job(func, types; kwargs...)
89+
GPUCompiler.cached_compilation(runtime_cache, job.source, job.config, compiler, linker)
90+
end
91+
7492
end

0 commit comments

Comments
 (0)