Add disk cache infrastructure for Julia 1.11

Project.toml

@@ -9,7 +9,10 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
+Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 

src/GPUCompiler.jl

@@ -9,7 +9,9 @@ using ExprTools: splitdef, combinedef
 
 using Libdl
 
+using Serialization
 using Scratch: @get_scratch!
+using Preferences
 
 const CC = Core.Compiler
 using Core: MethodInstance, CodeInstance, CodeInfo

src/execution.jl

@@ -61,6 +61,31 @@ end
 
 
 ## cached compilation
+disk_cache() = parse(Bool, @load_preference("disk_cache", "false"))
+
+"""
+    enable_cache!(state::Bool=true)
+
+Activate the GPUCompiler disk cache in the current environment.
+You will need to restart your Julia environment for it to take effect.
+
+!!! note
+    The cache functionality requires Julia 1.11
+"""
+function enable_cache!(state::Bool=true)
+    @set_preferences!("disk_cache"=>string(state))
+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()
 
@@ -115,19 +140,42 @@ end
 
     # fast path: find an applicable CodeInstance and see if we have compiled it before
     ci = ci_cache_lookup(ci_cache(job), src, world, world)::Union{Nothing,CodeInstance}
-    if ci !== nothing && haskey(cache, ci)
-        obj = cache[ci]
+    if ci !== nothing
+        obj = get(cache, ci, nothing)
     end
 
     # slow path: compile and link
     if obj === nothing || compile_hook[] !== nothing
-        # TODO: consider loading the assembly from an on-disk cache here
-        asm = compiler(job)
-
         if obj !== nothing
             # we got here because of a *compile* hook; don't bother linking
             return obj
         end
+        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())
+                end
+            end
+        else
+
+        if asm === nothing
+            asm = compiler(job)
+        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)
+        end
 
         obj = linker(job, asm)
         ci = ci_cache_lookup(ci_cache(job), src, world, world)::CodeInstance

test/CacheEnv/LocalPreferences.toml

@@ -0,0 +1,3 @@
+[GPUCompiler]
+disk_cache = "true"
+cache_key = "test"

test/CacheEnv/Project.toml

@@ -0,0 +1,2 @@
+[extras]
+GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"

test/cache.jl

@@ -0,0 +1,54 @@
+# @testset "Disk cache" begin
+#     @test GPUCompiler.disk_cache() == false
+#     cmd = Base.julia_cmd()
+#     if Base.JLOptions().project != C_NULL
+#         cmd = `$cmd --project=$(unsafe_string(Base.JLOptions().project))`
+#     end
+
+#     withenv("JULIA_LOAD_PATH" => "$(get(ENV, "JULIA_LOAD_PATH", "")):$(joinpath(@__DIR__, "CacheEnv"))") do
+#         @test success(pipeline(`$cmd cache.jl true`, stderr=stderr, stdout=stdout))
+#         @test success(pipeline(`$cmd cache.jl false`, stderr=stderr, stdout=stdout))
+#     end
+# end
+
+
+using GPUCompiler
+using Test
+
+const TOTAL_KERNELS = 1
+
+clear = parse(Bool, ARGS[1])
+
+@test GPUCompiler.disk_cache() == true
+
+if clear
+    GPUCompiler.clear_disk_cache!()
+    @test length(readdir(GPUCompiler.cache_path())) == 0
+else
+    @test length(readdir(GPUCompiler.cache_path())) == TOTAL_KERNELS
+end
+
+using LLVM, LLVM.Interop
+
+include("util.jl")
+include("definitions/native.jl")
+
+kernel() = return
+
+const runtime_cache = Dict{UInt, Any}()
+
+function compiler(job)
+    return GPUCompiler.compile(:asm, job)
+end
+
+function linker(job, asm)
+    asm
+end
+
+let (job, kwargs) = native_job(kernel, Tuple{})
+    source = job.source
+    config = job.config
+    GPUCompiler.cached_compilation(runtime_cache, config, source.ft, source.tt, compiler, linker)
+end
+
+@test length(readdir(GPUCompiler.cache_path())) == TOTAL_KERNELS