feat: preferences + docs

Author: avik-pal

docs/src/.vitepress/config.mts

@@ -89,6 +89,7 @@ export default defineConfig({
           {text: "Local build", link: "/tutorials/local-build"},
           {text: "Control Flow", link: "/tutorials/control-flow"},
           {text: "Sharding", link: "/tutorials/sharding"},
+          {text: "Persistent Compilation Cache", link: "/tutorials/persistent_compile_cache"},
         ],
       },
       {
@@ -160,6 +161,7 @@ export default defineConfig({
           { text: "Local build", link: "/tutorials/local-build" },
           { text: "Control Flow", link: "/tutorials/control-flow" },
           { text: "Sharding", link: "/tutorials/sharding" },
+          { text: "Persistent Compilation Cache", link: "/tutorials/persistent_compile_cache" },
         ],
       }
     ],

docs/src/tutorials/index.md

@@ -1,8 +1,10 @@
 # Tutorials
 
- - [Profiling](@ref profiling).
- - [Multi-Host Environments](@ref distributed).
- - [Local build of ReactantExtra](@ref local-build).
- - [Control flow](@ref control-flow).
+- [Profiling](@ref profiling).
+- [Multi-Host Environments](@ref distributed).
+- [Local build of ReactantExtra](@ref local-build).
+- [Control flow](@ref control-flow).
+- [Sharding](@ref sharding).
+- [Persistent Compilation Cache](@ref persistent_compile_cache).
 
 We are currently working on adding more tutorials to Reactant!! Please check back soon!

docs/src/tutorials/persistent_compile_cache.md

@@ -0,0 +1,28 @@
+# [Persistent Compilation Cache](@id persistent_compile_cache)
+
+Reactant.jl supports a persistent compilation cache that caches compiled and autotuned
+kernels on disk. We use [XLA's persisted autotuning](https://openxla.org/xla/persisted_autotuning)
+for this purpose. By default, the autotuning cache is enabled.
+
+## Preferences
+
+- `persistent_cache_enabled`: Whether to enable the persistent compilation cache. Defaults
+  to `false`.
+- `persistent_cache_directory`: The base directory to use for the persistent compilation
+  cache. Note that it is recommended to not set this preference, as Reactant will create
+  a unique directory corresponding to XLA and Reactant_jll's version. If the user sets
+  this preference, it is the user's responsibility to ensure that the directory exists
+  and is writable and needs to be segregated based on XLA and Reactant_jll's version.
+  Defaults to `""`.
+- `persistent_kernel_cache_enabled`: Whether to enable the kernel cache. Defaults to `false`.
+- `persistent_autotune_cache_enabled`: Whether to enable the autotuning cache. Defaults to
+  `true`.
+
+## Clearing the cache
+
+To clear the cache, you can reuse `Scratch.jl`'s `clear_scclear_scratchspaces!` function:
+
+```julia
+using Scratch, Reactant
+Scratch.clear_scratchspaces!(Reactant)
+```

src/PersistentCompileCache.jl

@@ -1,26 +1,45 @@
 module PersistentCompileCache
 
+using ..Reactant: Reactant
+
+using Preferences: load_preference
 using Scratch: @get_scratch!
 using Reactant_jll: Reactant_jll
 
 const CACHE_DIR = Ref{Union{Nothing,String}}(nothing)
 const KERNEL_CACHE_ENABLED = Ref(false)
 const AUTOTUNE_CACHE_ENABLED = Ref(false)
 
-# TODO: preference to disable persistent caching
-# TODO: preference to control what is being cached.
-# TODO: preference for cache dir
-# TODO: clear cache function
-
 function __init__()
-    # We version our cache directory based on Reactant_jll version (technically we need to
-    # version according to XLA, but this is a good enough proxy)
-    version = pkgversion(Reactant_jll)
-    CACHE_DIR[] = @get_scratch!(
-        "xla_persistent_cache_$(version.major)_$(version.minor)_$(version.patch)"
-    )
-    KERNEL_CACHE_ENABLED[] = true
-    AUTOTUNE_CACHE_ENABLED[] = true
+    persistent_cache_enabled = load_preference(Reactant, "persistent_cache_enabled", true)
+    persistent_cache_directory = load_preference(Reactant, "persistent_cache_directory", "")
+
+    if persistent_cache_enabled
+        if isempty(persistent_cache_directory)
+            # We version our cache directory based on Reactant_jll version (technically we
+            # need to version according to XLA, but this is a good enough proxy)
+            version = pkgversion(Reactant_jll)
+            CACHE_DIR[] = @get_scratch!(
+                "xla_persistent_cache_$(version.major)_$(version.minor)_$(version.patch)"
+            )
+        else
+            CACHE_DIR[] = persistent_cache_directory
+        end
+        @debug "Persistent compilation cache enabled. Using base directory: $(CACHE_DIR[])"
+
+        KERNEL_CACHE_ENABLED[] = load_preference(
+            Reactant, "persistent_kernel_cache_enabled", false
+        )
+        @debug "Kernel cache enabled: $(KERNEL_CACHE_ENABLED[])"
+
+        AUTOTUNE_CACHE_ENABLED[] = load_preference(
+            Reactant, "persistent_autotune_cache_enabled", true
+        )
+        @debug "Autotune cache enabled: $(AUTOTUNE_CACHE_ENABLED[])"
+    else
+        @debug "Persistent compilation cache disabled..."
+    end
+
     return nothing
 end
 

test/basic.jl

@@ -1510,3 +1510,17 @@ end
         @test @jit(nested_mapreduce_hcat(x_ra, y_ra)) ≈ nested_mapreduce_hcat(x, y)
     end
 end
+
+@testset "compilation cache" begin
+    if Reactant.PersistentCompileCache.autotune_cache_enabled() &&
+        contains(string(Reactant.devices()[1]), "CUDA")
+        A = Reactant.to_rarray(rand(Float32, 2, 5))
+        B = Reactant.to_rarray(rand(Float32, 5, 1000))
+        @jit A * B # This should populate the cache dir
+
+        @test any(
+            endswith(".textproto"),
+            readdir(Reactant.PersistentCompileCache.get_autotune_cache_directory()),
+        )
+    end
+end