From 749bec2a2569595f6158bb99c5d32a3b1a600a5b Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 15 Jan 2025 12:01:22 -0600
Subject: [PATCH] Generalize precompilation support

---
 src/Precompile.jl | 35 ++++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/Precompile.jl b/src/Precompile.jl
index 98c60dee5..ce692909a 100644
--- a/src/Precompile.jl
+++ b/src/Precompile.jl
@@ -34,15 +34,29 @@ function infer_sig(sig)
     end
 end
 
+function clear_oc_cache()
+    # Opaque closures capture the worldage of their compilation and thus are not relocatable
+    # Therefore we explicitly purge all OC's we have created here
+    for v in oc_capture_vec
+        if v isa Base.RefValue
+            p = Ptr{Ptr{Cvoid}}(pointer_from_objref(v))
+            Base.atomic_pointerset(p, C_NULL, :monotonic)
+        else
+            empty!(v)
+        end
+    end
+end
+
+# Precompilation on 1.10 hits an apparent bug: https://github.com/JuliaLang/julia/issues/56947
+function precompilation_supported()
+    return VERSION >= v"1.11" || VERSION >= v"1.10.8"
+end
+
 @setup_workload begin
     initialize_dialect()
     client = XLA.CPUClient(; checkcount=false)
     @compile_workload begin
-        # Precompilation on 1.10 hits an apparent bug: https://github.com/JuliaLang/julia/issues/56947
-        @static if VERSION < v"1.11"
-        else
-            # infer_sig(Tuple{typeof(Base.sum), Reactant.TracedRArray{Float64, 2}})
-            # infer_sig(Tuple{typeof(Base.sin), Reactant.TracedRNumber{Float64}})
+	@static if precompilation_supported()
             x = ConcreteRNumber(2.0; client)
             Reactant.compile(sin, (x,); client)
 
@@ -53,14 +67,5 @@ end
     XLA.free_client(client)
     client.client = C_NULL
     deinitialize_dialect()
-    # Opaque closures capture the worldage of their compilation and thus are not relocatable
-    # Therefore we explicitly purge all OC's we have created here
-    for v in oc_capture_vec
-        if v isa Base.RefValue
-            p = Ptr{Ptr{Cvoid}}(pointer_from_objref(v))
-            Base.atomic_pointerset(p, C_NULL, :monotonic)
-        else
-            empty!(v)
-        end
-    end
+    clear_oc_cache()
 end