diff --git a/Project.toml b/Project.toml index 0a8ef9e338..7a2040f9ef 100644 --- a/Project.toml +++ b/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" diff --git a/src/Enzyme.jl b/src/Enzyme.jl index 942df0581c..92ec9a623f 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -1587,4 +1587,6 @@ Returns true if within autodiff, otherwise false. """ @inline EnzymeCore.within_autodiff() = false +include("precompile.jl") + end # module diff --git a/src/compiler/orcv2.jl b/src/compiler/orcv2.jl index 7588eddb78..1640b05db2 100644 --- a/src/compiler/orcv2.jl +++ b/src/compiler/orcv2.jl @@ -83,7 +83,7 @@ function define_absolute_symbol(jd, name) return false end -function __init__() +function setup_globals() opt_level = Base.JLOptions().opt_level if opt_level < 2 optlevel = LLVM.API.LLVMCodeGenLevelNone @@ -105,11 +105,6 @@ function __init__() dg = LLVM.CreateDynamicLibrarySearchGeneratorForProcess(prefix) LLVM.add!(jd_main, dg) - if Sys.iswindows() && Int === Int64 - # TODO can we check isGNU? - define_absolute_symbol(jd_main, mangle(lljit, "___chkstk_ms")) - end - es = ExecutionSession(lljit) try lctm = LLVM.LocalLazyCallThroughManager(triple(lljit), es) @@ -120,6 +115,17 @@ function __init__() jit[] = CompilerInstance(lljit, nothing, nothing) end + jd_main, lljit +end + +function __init__() + jd_main, lljit = setup_globals() + + if Sys.iswindows() && Int === Int64 + # TODO can we check isGNU? + define_absolute_symbol(jd_main, mangle(lljit, "___chkstk_ms")) + end + hnd = unsafe_load(cglobal(:jl_libjulia_handle, Ptr{Cvoid})) for (k, v) in Compiler.JuliaGlobalNameMap ptr = unsafe_load(Base.reinterpret(Ptr{Ptr{Cvoid}}, Libdl.dlsym(hnd, k))) diff --git a/src/precompile.jl b/src/precompile.jl new file mode 100644 index 0000000000..c20eaac149 --- /dev/null +++ b/src/precompile.jl @@ -0,0 +1,13 @@ +using PrecompileTools: @setup_workload, @compile_workload + +@setup_workload begin + precompile_module = @eval module $(gensym()) + f(x) = x^2 + end + + Compiler.JIT.setup_globals() + + @compile_workload begin + Enzyme.autodiff(Reverse, precompile_module.f, Active(2.0)) + end +end