From 33171a12fce924df5265f887e7f0f01a6f626bcd Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 2 Aug 2022 10:54:08 -0400
Subject: [PATCH 01/11] Add disk cache infrastructure back with tests

Apply suggestions from code review
---
 Manifest.toml                       | 20 ++++--------
 Project.toml                        |  3 ++
 src/GPUCompiler.jl                  |  2 ++
 src/cache.jl                        | 50 ++++++++++++++++++++++++++++-
 test/CacheEnv/LocalPreferences.toml |  3 ++
 test/CacheEnv/Project.toml          |  2 ++
 test/cache.jl                       | 38 ++++++++++++++++++++++
 test/runtests.jl                    | 14 ++++++++
 8 files changed, 118 insertions(+), 14 deletions(-)
 create mode 100644 test/CacheEnv/LocalPreferences.toml
 create mode 100644 test/CacheEnv/Project.toml
 create mode 100644 test/cache.jl

diff --git a/Manifest.toml b/Manifest.toml
index b71ce277..01dcd8c0 100644
--- a/Manifest.toml
+++ b/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"
 
diff --git a/Project.toml b/Project.toml
index 5b6c17a7..d7505124 100644
--- a/Project.toml
+++ b/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"
 
diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl
index 34f3fbd6..b5c68ea0 100644
--- a/src/GPUCompiler.jl
+++ b/src/GPUCompiler.jl
@@ -9,7 +9,9 @@ using ExprTools: splitdef, combinedef
 
 using Libdl
 
+using Serialization
 using Scratch: @get_scratch!
+using Preferences
 
 include("utils.jl")
 
diff --git a/src/cache.jl b/src/cache.jl
index 32286a9f..04ae1492 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -127,6 +127,39 @@ end
                 true)))
 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 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) * "-" * key(pkg_version))
+clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
+
 const cache_lock = ReentrantLock()
 
 """
@@ -173,7 +206,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 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
@@ -182,6 +226,10 @@ end
         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,
diff --git a/test/CacheEnv/LocalPreferences.toml b/test/CacheEnv/LocalPreferences.toml
new file mode 100644
index 00000000..1312728b
--- /dev/null
+++ b/test/CacheEnv/LocalPreferences.toml
@@ -0,0 +1,3 @@
+[GPUCompiler]
+disk_cache = true
+cache_key = "test"
diff --git a/test/CacheEnv/Project.toml b/test/CacheEnv/Project.toml
new file mode 100644
index 00000000..092c0f04
--- /dev/null
+++ b/test/CacheEnv/Project.toml
@@ -0,0 +1,2 @@
+[extras]
+GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
diff --git a/test/cache.jl b/test/cache.jl
new file mode 100644
index 00000000..78cf5ef6
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index bb8aec55..c703fee0 100644
--- a/test/runtests.jl
+++ b/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

From da0edc6c90e77d99e81dc0dc8ce808ff5f87dd7a Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Wed, 22 Mar 2023 23:04:45 -0400
Subject: [PATCH 02/11] fix pkg_version

---
 src/cache.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index 04ae1492..71809e5c 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -157,7 +157,7 @@ function set_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))
+cache_path() = @get_scratch!(cache_key * "-kernels-" * key(VERSION) * "-" * key(Base.pkgversion(@__MODULE__)))
 clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
 
 const cache_lock = ReentrantLock()

From 939bf4e711753876a5b49c8d0e1b7a5ed18055b7 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 26 Mar 2023 19:01:49 -0400
Subject: [PATCH 03/11] Make caching runtime not compiletime

---
 src/cache.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/cache.jl b/src/cache.jl
index 71809e5c..445667cc 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -127,8 +127,8 @@ end
                 true)))
 end
 
-const disk_cache = parse(Bool, @load_preference("disk_cache", "false"))
-const cache_key = @load_preference("cache_key", "")
+disk_cache() = parse(Bool, @load_preference("disk_cache", "false"))
+cache_key() = @load_preference("cache_key", "")
 
 """
     enable_cache!(state=true)
@@ -157,7 +157,7 @@ function set_cache_key(key)
 end
 
 key(ver::VersionNumber) = "$(ver.major)_$(ver.minor)_$(ver.patch)"
-cache_path() = @get_scratch!(cache_key * "-kernels-" * key(VERSION) * "-" * key(Base.pkgversion(@__MODULE__)))
+cache_path() = @get_scratch!(cache_key() * "-kernels-" * key(VERSION) * "-" * key(Base.pkgversion(@__MODULE__)))
 clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
 
 const cache_lock = ReentrantLock()
@@ -207,7 +207,7 @@ end
 
     asm = nothing
     # can we load from the disk cache?
-    if disk_cache
+    if disk_cache()
         path = joinpath(cache_path(), "$key.jls")
         if isfile(path)
             try

From 639ac3047fdb8074144f8e58db0fd7f7f1f45232 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Mon, 27 Mar 2023 12:15:11 -0400
Subject: [PATCH 04/11] fixup! Make caching runtime not compiletime

---
 src/cache.jl     | 2 +-
 test/cache.jl    | 2 +-
 test/runtests.jl | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/cache.jl b/src/cache.jl
index 445667cc..bd77f72a 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -227,7 +227,7 @@ end
 
         asm = compiler(job)
 
-        if disk_cache && !isfile(path)
+        if disk_cache() && !isfile(path)
             serialize(path, asm)
         end
     end
diff --git a/test/cache.jl b/test/cache.jl
index 78cf5ef6..e37adea3 100644
--- a/test/cache.jl
+++ b/test/cache.jl
@@ -5,7 +5,7 @@ const TOTAL_KERNELS = 1
 
 clear = parse(Bool, ARGS[1])
 
-@test GPUCompiler.disk_cache == true
+@test GPUCompiler.disk_cache() == true
 
 if clear
     GPUCompiler.clear_disk_cache!()
diff --git a/test/runtests.jl b/test/runtests.jl
index c703fee0..66db331d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -30,7 +30,7 @@ include("examples.jl")
 haskey(ENV, "CI") && GPUCompiler.timings()
 
 @testset "Disk cache" begin
-    @test GPUCompiler.disk_cache == false
+    @test GPUCompiler.disk_cache() == false
 
     cmd = Base.julia_cmd()
     if Base.JLOptions().project != C_NULL

From c585f72e262cb3545d8b3e5db4a638eca5815cdf Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Mon, 27 Mar 2023 15:56:09 -0400
Subject: [PATCH 05/11] fixup! fixup! Make caching runtime not compiletime

---
 src/cache.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index bd77f72a..e0ce60dd 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -157,7 +157,7 @@ function set_cache_key(key)
 end
 
 key(ver::VersionNumber) = "$(ver.major)_$(ver.minor)_$(ver.patch)"
-cache_path() = @get_scratch!(cache_key() * "-kernels-" * key(VERSION) * "-" * key(Base.pkgversion(@__MODULE__)))
+cache_path() = @get_scratch!(cache_key() * "-kernels-" * key(VERSION))
 clear_disk_cache!() = rm(cache_path(); recursive=true, force=true)
 
 const cache_lock = ReentrantLock()

From 5d2d446a4f844ea7a6b6454f93d233b0d87cd1f8 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 11 Apr 2023 19:49:27 -0400
Subject: [PATCH 06/11] add environemnt

---
 src/cache.jl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index e0ce60dd..3a7b8c40 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -128,7 +128,11 @@ end
 end
 
 disk_cache() = parse(Bool, @load_preference("disk_cache", "false"))
-cache_key() = @load_preference("cache_key", "")
+function cache_key() =
+    major = @load_preference("cache_key", "")
+    minor = get(ENV, "JULIA_GPUCOMPILER_CACHE", "")
+    string(major, "-" minor)
+end
 
 """
     enable_cache!(state=true)

From a7a5a717a1d2eb4b22ef16ee4209c20e5c8d1502 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 11 Apr 2023 19:51:19 -0400
Subject: [PATCH 07/11] fixup! add environemnt

---
 src/cache.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index 3a7b8c40..081525c1 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -128,7 +128,7 @@ end
 end
 
 disk_cache() = parse(Bool, @load_preference("disk_cache", "false"))
-function cache_key() =
+function cache_key()
     major = @load_preference("cache_key", "")
     minor = get(ENV, "JULIA_GPUCOMPILER_CACHE", "")
     string(major, "-" minor)

From ca8e8bbdca046fde76fea9ed08a2beb7a11c4371 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 11 Apr 2023 19:53:40 -0400
Subject: [PATCH 08/11] fixup! add environemnt

---
 src/cache.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index 081525c1..a5878936 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -131,7 +131,7 @@ 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)
+    string(major, "-", minor)
 end
 
 """

From c8b4f9a69e688720e2efc2dad02cf370a92f6659 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 11 Apr 2023 21:04:29 -0400
Subject: [PATCH 09/11] fix debug statement

---
 src/cache.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index a5878936..7823da71 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -216,7 +216,7 @@ end
         if isfile(path)
             try
                 asm = deserialize(path)
-                @debug "Loading compiled kernel for $spec from $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

From c95966f0aaeb4a69f2a8f3bc840be4606b57d416 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Wed, 12 Apr 2023 14:24:29 -0400
Subject: [PATCH 10/11] fix race condition observed at scale

---
 src/cache.jl                        | 14 ++++++++++----
 test/CacheEnv/LocalPreferences.toml |  2 +-
 test/cache.jl                       |  6 ++++--
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/cache.jl b/src/cache.jl
index 7823da71..fb03e6d6 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -135,7 +135,7 @@ function cache_key()
 end
 
 """
-    enable_cache!(state=true)
+    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.
@@ -145,8 +145,8 @@ You will need to restart your Julia environment for it to take effect.
     `cache_key` (see [`set_cache_key``](@ref)), the GPUCompiler version
     and your Julia version.
 """
-function enable_cache!(state=true)
-    @set_preferences!("disk_cache"=>state)
+function enable_cache!(state::Bool=true)
+    @set_preferences!("disk_cache"=>string(state))
 end
 
 """
@@ -232,7 +232,13 @@ end
         asm = compiler(job)
 
         if disk_cache() && !isfile(path)
-            serialize(path, asm)
+            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
+            mv(tmppath, path, force=true)
         end
     end
 
diff --git a/test/CacheEnv/LocalPreferences.toml b/test/CacheEnv/LocalPreferences.toml
index 1312728b..4509f744 100644
--- a/test/CacheEnv/LocalPreferences.toml
+++ b/test/CacheEnv/LocalPreferences.toml
@@ -1,3 +1,3 @@
 [GPUCompiler]
-disk_cache = true
+disk_cache = "true"
 cache_key = "test"
diff --git a/test/cache.jl b/test/cache.jl
index e37adea3..2bbcb63f 100644
--- a/test/cache.jl
+++ b/test/cache.jl
@@ -32,7 +32,9 @@ function linker(job, asm)
 end
 
 let (job, kwargs) = native_job(kernel, Tuple{})
-    GPUCompiler.cached_compilation(runtime_cache, job, compiler, linker)
+    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
\ No newline at end of file
+@test length(readdir(GPUCompiler.cache_path())) == TOTAL_KERNELS

From 051e7955a2302841d81b166acc7423fd94ff07b1 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Wed, 12 Apr 2023 15:38:03 -0400
Subject: [PATCH 11/11] fixup! fix race condition observed at scale

---
 src/cache.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cache.jl b/src/cache.jl
index fb03e6d6..c3e064cd 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -238,7 +238,7 @@ end
             serialize(io, asm)
             close(io)
             # atomic move
-            mv(tmppath, path, force=true)
+            Base.rename(tmppath, path, force=true)
         end
     end