Add disk cache infrastructure back with tests

Manifest.toml

@@ -23,9 +23,9 @@ deps = ["ArgTools", "LibCURL", "NetworkOptions"]
 uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 
 [[ExprTools]]
-git-tree-sha1 = "c1d06d129da9f55715c6c212866f5b1bddc5fa00"
+git-tree-sha1 = "56559bbef6ca5ea0c0818fa5c90320398a6fbf8d"
 uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
-version = "0.1.9"
+version = "0.1.8"
 
 [[InteractiveUtils]]
 deps = ["Markdown"]
@@ -39,15 +39,15 @@ version = "1.4.1"
 
 [[LLVM]]
 deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
-git-tree-sha1 = "1c614dfbecbaee4897b506bba2b432bf0d21f2ed"
+git-tree-sha1 = "df115c31f5c163697eede495918d8e85045c8f04"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
-version = "4.17.0"
+version = "4.16.0"
 
 [[LLVMExtra_jll]]
-deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
-git-tree-sha1 = "e46e3a40daddcbe851f86db0ec4a4a3d4badf800"
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]
+git-tree-sha1 = "7718cf44439c676bc0ec66a87099f41015a522d6"
 uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
-version = "0.0.19+0"
+version = "0.0.16+2"
 
 [[LazyArtifacts]]
 deps = ["Artifacts", "Pkg"]
@@ -114,12 +114,6 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 [[SHA]]
 uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 
-[[Scratch]]
-deps = ["Dates"]
-git-tree-sha1 = "30449ee12237627992a99d5e30ae63e4d78cd24a"
-uuid = "6c6a2e73-6563-6170-7368-637461726353"
-version = "1.2.0"
-
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 

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
 
 include("utils.jl")
 

src/cache.jl

@@ -127,6 +127,43 @@ end
                 true)))
 end
 
+disk_cache() = parse(Bool, @load_preference("disk_cache", "false"))
+function cache_key()
+    major = @load_preference("cache_key", "")
+    minor = get(ENV, "JULIA_GPUCOMPILER_CACHE", "")
+    string(major, "-", minor)
+end
+
+"""
+    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.
+
+!!! 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::Bool=true)
+    @set_preferences!("disk_cache"=>string(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 spurious 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))
+clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
+
 const cache_lock = ReentrantLock()
 
 """
@@ -173,7 +210,18 @@ end
     job = CompilerJob(src, cfg)
 
     asm = nothing
-    # TODO: consider loading the assembly from an on-disk cache here
+    # can we load from the disk cache?
+    if disk_cache()
+        path = joinpath(cache_path(), "$key.jls")
+        if isfile(path)
+            try
+                asm = deserialize(path)
+                @debug "Loading compiled kernel from $path" ft tt world cfg
+            catch ex
+                @warn "Failed to load compiled kernel at $path" exception=(ex, catch_backtrace())
+            end
+        end
+    end
 
     # compile
     if asm === nothing
@@ -182,6 +230,16 @@ end
         end
 
         asm = compiler(job)
+
+        if disk_cache() && !isfile(path)
+            tmppath, io = mktemp(;cleanup=false)
+            # TODO: We should add correctness checks.
+            #       Like size of data, as well as ft, tt, world, cfg
+            serialize(io, asm)
+            close(io)
+            # atomic move
+            Base.rename(tmppath, path, force=true)
+        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,40 @@
+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

test/runtests.jl

@@ -29,4 +29,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