Add disk cache infrastructure for Julia 1.11

Project.toml

@@ -9,17 +9,23 @@ 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"
 
 [compat]
 ExprTools = "0.1"
 InteractiveUtils = "1"
+LLVM = "6.6"
 Libdl = "1"
 Logging = "1"
-UUIDs = "1"
-LLVM = "6.6"
+Preferences = "1"
 Scratch = "1"
+Serialization = "1"
+TOML = "1"
 TimerOutputs = "0.5"
+UUIDs = "1"
 julia = "1.8"

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,23 @@ 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)
 
 const cache_lock = ReentrantLock()
 
@@ -108,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)
@@ -117,20 +158,53 @@ end
     ci = ci_cache_lookup(ci_cache(job), src, world, world)::Union{Nothing,CodeInstance}
     if ci !== nothing
         key = (ci, cfg)
-        if haskey(cache, key)
-            obj = cache[key]
-        end
+        obj = get(cache, key, 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)
+        asm = nothing
+        path = nothing
+        ondisk_hit = false
+        @static if VERSION >= v"1.11.0-"
+            # Don't try to hit the disk cache if we are for a *compile* hook
+            if ci !== nothing && obj === nothing && disk_cache() # TODO: (Should we allow backends to opt out?)
+                path = cache_file(ci, cfg)
+                @debug "Looking for on-disk cache" job path
+                if path !== nothing && isfile(path)
+                    ondisk_hit = true
+                    try
+                        @debug "Loading compiled kernel" job path
+                        asm = deserialize(path)
+                    catch ex
+                        @warn "Failed to load compiled kernel" job path exception=(ex, catch_backtrace())
+                    end
+                end
+            end
+        end
 
+        if asm === nothing || compile_hook[] !== nothing
+            # Run the compiler in-case we need to hook it.
+            asm = compiler(job)
+        end
         if obj !== nothing
             # we got here because of a *compile* hook; don't bother linking
             return obj
         end
+
+        @static if VERSION >= v"1.11.0-"
+            if !ondisk_hit && path !== nothing && disk_cache()
+                @debug "Writing out on-disk cache" job path
+                # 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
+                mkpath(dirname(path))
+                Base.rename(tmppath, path, force=true)
+            end
+        end
+
         obj = linker(job, asm)
 
         if ci === nothing

test/native_tests.jl

@@ -549,12 +549,13 @@ precompile_test_harness("Inference caching") do load_path
         import GPUCompiler
         using PrecompileTools
 
-        function kernel()
+        function kernel(A, x)
+            A[1] = x
             return
         end
 
         let
-            job, _ = NativeCompiler.create_job(kernel, ())
+            job, _ = NativeCompiler.create_job(kernel, (Vector{Int}, Int))
             GPUCompiler.code_typed(job)
         end
 
@@ -578,20 +579,30 @@ precompile_test_harness("Inference caching") do load_path
             job, _ = NativeCompiler.create_job(identity, (Int,))
             GPUCompiler.ci_cache_token(job)
         end
-        ci = isdefined(identity_mi, :cache) ? identity_mi.cache : nothing
-        while ci !== nothing
-            @test ci.owner !== token
-            ci = isdefined(ci, :next) ? ci.next : nothing
-        end
+        @test !check_presence(identity_mi, token)
 
         using InferenceCaching
 
         # Check that kernel survived
-        kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{})
+        kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{Vector{Int}, Int})
         @test check_presence(kernel_mi, token)
 
         # check that identity survived
         @test check_presence(identity_mi, token)
+
+        GPUCompiler.enable_cache!()
+        @test GPUCompiler.disk_cache() == true
+
+        job, _ = NativeCompiler.create_job(InferenceCaching.kernel, (Vector{Int}, Int))
+        @assert job.source == kernel_mi
+        ci = GPUCompiler.ci_cache_lookup(GPUCompiler.ci_cache(job), job.source, job.world, job.world)
+        @assert ci !== nothing
+        @assert ci.inferred !== nothing
+        path = GPUCompiler.cache_file(ci, job.config)
+        @test path !== nothing
+        @test !ispath(path)
+        NativeCompiler.cached_execution(InferenceCaching.kernel, (Vector{Int}, Int))
+        @test ispath(path)
     end
 end
 

test/native_testsetup.jl

@@ -71,4 +71,22 @@ function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
     end
 end
 
+const runtime_cache = Dict{Any, Any}()
+
+function compiler(job)
+    JuliaContext() do ctx
+        GPUCompiler.compile(:asm, job, validate=false)
+    end
+end
+
+function linker(job, asm)
+    asm
+end
+
+# simulates cached codegen
+function cached_execution(@nospecialize(func), @nospecialize(types); kwargs...)
+    job, kwargs = create_job(func, types; kwargs...)
+    GPUCompiler.cached_compilation(runtime_cache, job.source, job.config, compiler, linker)
+end
+
 end