Skip to content

Commit

Permalink
Generalize precompilation support
Browse files Browse the repository at this point in the history
  • Loading branch information
wsmoses committed Jan 15, 2025
1 parent 4b9434e commit 749bec2
Showing 1 changed file with 20 additions and 15 deletions.
35 changes: 20 additions & 15 deletions src/Precompile.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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

0 comments on commit 749bec2

Please sign in to comment.