Add tests for precompilation support

Author: vchuravy

Project.toml

@@ -12,6 +12,7 @@ LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ObjectFile = "d8793406-e978-5875-9003-1fc021f44a92"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -34,6 +35,7 @@ Enzyme_jll = "0.0.123"
 GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26"
 LLVM = "6.1, 7"
 ObjectFile = "0.4"
+PrecompileTools = "1.2"
 Preferences = "1.4"
 SpecialFunctions = "1, 2"
 StaticArrays = "1"

test/Project.toml

@@ -12,6 +12,7 @@ LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LLVM_jll = "86de99a1-58d6-5da7-8064-bd56ce2e322c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

test/precompile.jl

@@ -0,0 +1,59 @@
+using Test
+
+function precompile_test_harness(@nospecialize(f), testset::String)
+    @testset "$testset" begin
+        precompile_test_harness(f, true)
+    end
+end
+function precompile_test_harness(@nospecialize(f), separate::Bool)
+    load_path = mktempdir()
+    load_cache_path = separate ? mktempdir() : load_path
+    try
+        pushfirst!(LOAD_PATH, load_path)
+        pushfirst!(DEPOT_PATH, load_cache_path)
+        f(load_path)
+    finally
+        try
+            rm(load_path, force=true, recursive=true)
+        catch err
+            @show err
+        end
+        if separate
+            try
+                rm(load_cache_path, force=true, recursive=true)
+            catch err
+                @show err
+            end
+        end
+        filter!((≠)(load_path), LOAD_PATH)
+        separate && filter!((≠)(load_cache_path), DEPOT_PATH)
+    end
+    nothing
+end
+
+precompile_test_harness("Inference caching") do load_path
+    write(joinpath(load_path, "InferenceCaching.jl"), :(module InferenceCaching
+        using Enzyme
+        using PrecompileTools
+
+        function mul(x, y)
+            return x * y
+        end
+
+        @setup_workload begin
+            @compile_workload begin
+                autodiff(Reverse, mul, Active, Active(1.0), Active(2.0))
+                autodiff(Forward, mul, Duplicated, Duplicated(1.0, 1.0), Const(2.0))
+            end
+        end
+    end) |> string)
+
+    Base.compilecache(Base.PkgId("InferenceCaching"))
+    @eval let
+        using InferenceCaching
+        using Enzyme
+
+        autodiff(Reverse, InferenceCaching.mul, Active, Active(1.0), Active(2.0))
+        autodiff(Forward, InferenceCaching.mul, Duplicated, Duplicated(1.0, 1.0), Const(2.0))
+    end
+end