Add disk cache infrastructure back with tests

Manifest.toml

@@ -114,6 +114,12 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 [[SHA]]
 uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 
+[[Scratch]]
+deps = ["Dates"]
+git-tree-sha1 = "f94f779c94e58bf9ea243e77a37e16d9de9126bd"
+uuid = "6c6a2e73-6563-6170-7368-637461726353"
+version = "1.1.1"
+
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 

Project.toml

@@ -9,6 +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,6 +9,17 @@ using ExprTools: splitdef, combinedef
 
 using Libdl
 
+using Preferences
+using Scratch
+using Serialization
+
+using TOML
+# Get the current version at compile-time, that's fine it's not going to change. ;)
+function get_version()
+    return VersionNumber(TOML.parsefile(joinpath(dirname(@__DIR__), "Project.toml"))["version"])
+end
+const pkg_version = get_version()
+
 const to = TimerOutput()
 
 timings() = (TimerOutputs.print_timer(to); println())

src/cache.jl

@@ -62,6 +62,40 @@ const specialization_counter = Ref{UInt}(0)
     return new_ci
 end
 
+const disk_cache = parse(Bool, @load_preference("disk_cache", "false"))
+const cache_key = @load_preference("cache_key", "")
+
+"""
+    enable_cache!(state=true)
+
+Activate the GPUCompiler disk cache in the current environment.
+You will need to restart your Julia environment for it to take effect.
+
+!!! warning
+    The disk cache is not automatically invalidated. It is sharded upon
+    `cache_key` (see [`set_cache_key``](@ref)), the GPUCompiler version
+    and your Julia version.
+"""
+function enable_cache!(state=true)
+    @set_preferences!("disk_cache"=>state)
+end
+
+"""
+    set_cache_key(key)
+
+If you are deploying an application it is recommended that you use your
+application name and version as a cache key. To minimize the risk of
+encountering spurios cache hits.
+"""
+function set_cache_key(key)
+    @set_preferences!("cache_key"=>key)
+end
+
+key(ver::VersionNumber) = "$(ver.major)_$(ver.minor)_$(ver.patch)"
+cache_path() = @get_scratch!(cache_key * "-kernels-" * key(VERSION) * "-" * key(pkg_version))
+clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
+
+
 const cache_lock = ReentrantLock()
 function cached_compilation(cache::AbstractDict,
                             @nospecialize(job::CompilerJob),
@@ -81,13 +115,30 @@ function cached_compilation(cache::AbstractDict,
         if obj === nothing || force_compilation
             asm = nothing
 
+            # can we load from the disk cache?
+            if disk_cache && !force_compilation
+                path = joinpath(cache_path(), "$key.jls")
+                if isfile(path)
+                    try
+                        asm = deserialize(path)
+                        @debug "Loading compiled kernel for $spec from $path"
+                    catch ex
+                        @warn "Failed to load compiled kernel at $path" exception=(ex, catch_backtrace())
+                    end
+                end
+            end
+
             # compile
             if asm === nothing
                 if compile_hook[] !== nothing
                     compile_hook[](job)
                 end
 
                 asm = compiler(job)
+
+                if disk_cache && !isfile(path)
+                    serialize(path, asm)
+                end
             end
 
             # link (but not if we got here because of forced compilation)

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,38 @@
+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{})
+    GPUCompiler.cached_compilation(runtime_cache, job, compiler, linker)
+end
+
+@test length(readdir(GPUCompiler.cache_path())) == TOTAL_KERNELS

test/runtests.jl

@@ -28,4 +28,18 @@ include("examples.jl")
 
 haskey(ENV, "CI") && GPUCompiler.timings()
 
+@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
+
 end