From 1c7641105f7cf9aaa58f2c07bf168623674adae4 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 13:13:05 +0100 Subject: [PATCH] Precompile the runtime when a GPU can be used while precompiling Allocating, copying, launching kernels and displaying arrays were not covered by the precompilation workload, which only compiles a kernel without a device, so they compiled seconds of host code on first use: on an Arc A770, the first array took 2.9 s, the first broadcast another 2.4 s. Exercise those operations as well when the GPU can be used while precompiling, and reset the state they create afterwards, as it refers to driver objects of the precompilation process. Precompilation doesn't initialize the driver when loading oneAPI, so add `can_precompile`, which initializes it (quietly, once) and checks that the runtime can be used: the SPIR-V tool for the codegen path, a device, and in-order immediate command lists. It can be disabled with the `precompile` preference, and is meant for the workloads of other packages too. What it finds is not kept in the package image. On Julia 1.13, this takes the first array to 13 ms, the first broadcast to 0.17 s and the first display of an array to 7 ms, for about 40 ms more load time and a package image that is 9 MB larger. On Julia 1.12, loading takes 65 to 110 ms longer, about 50 ms of which is loading the package image: that is the code for compiling and launching kernels the workload caches, and leaving out any one of its operations doesn't reduce it measurably. --- docs/src/installation.md | 21 ++++++++++++ lib/level-zero/oneL0.jl | 26 +++++++++------ src/compiler/precompile.jl | 66 +++++++++++++++++++++++++++++++++++++- src/oneAPI.jl | 41 ++++++++++++++++++++++- test/level-zero.jl | 8 +++++ 5 files changed, 151 insertions(+), 11 deletions(-) diff --git a/docs/src/installation.md b/docs/src/installation.md index 4f60a13a3..5da810db2 100644 --- a/docs/src/installation.md +++ b/docs/src/installation.md @@ -60,6 +60,27 @@ LLVM's SPIR-V back-end. See [Intel LTS Driver Stack](lts.md) for what else chang it costs. The switch does not install an LTS driver — combine it with the system-library configuration below. +## Precompilation + +When a GPU can be used while oneAPI.jl is precompiled, the precompilation workload also +allocates memory and launches kernels on it, so that doing so for the first time in a +session doesn't have to compile host code. Without a usable GPU, e.g. on a login node, +precompilation leaves that out. To leave it out regardless, disable the `precompile` +preference before oneAPI.jl is loaded (loading it may precompile it), e.g. in +`LocalPreferences.toml`: + +```toml +[oneAPI] +precompile = false +``` + +or with Preferences.jl, by oneAPI.jl's UUID: + +```julia +using Preferences, UUIDs +set_preferences!(UUID("8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"), "precompile" => false) +``` + ## Using System Libraries (Advanced) !!! warning diff --git a/lib/level-zero/oneL0.jl b/lib/level-zero/oneL0.jl index 565391f32..d654abd72 100644 --- a/lib/level-zero/oneL0.jl +++ b/lib/level-zero/oneL0.jl @@ -164,8 +164,16 @@ function parse_env_bool(name::AbstractString, default::Bool) end function __init__() + # precompilation processes only initialize the driver when a workload asks for it (see + # `oneAPI.can_precompile`) precompiling = ccall(:jl_generating_output, Cint, ()) != 0 precompiling && return + initialize() +end + +# initialize the driver, setting `functional`. `quiet` doesn't report why it isn't functional. +function initialize(; quiet::Bool=false) + functional[] = false # Driver handles are process-local, so nothing keyed on one may survive into another # process through a precompiled image. @@ -179,18 +187,18 @@ function __init__() if Sys.iswindows() if Libdl.dlopen(libze_loader; throw_error=false) === nothing - @error "The oneAPI Level Zero loader was not found. Please ensure the Intel GPU drivers are installed." + quiet || @error "The oneAPI Level Zero loader was not found. Please ensure the Intel GPU drivers are installed." return end else if !oneAPI_Level_Zero_Loader_jll.is_available() - @error """No oneAPI Level Zero loader found for your platform. Currently, only Linux x86 is supported. + quiet || @error """No oneAPI Level Zero loader found for your platform. Currently, only Linux x86 is supported. If you have a local oneAPI toolchain, you can use that; refer to the documentation for more details.""" return end if !NEO_jll.is_available() - @error """No oneAPI driver found for your platform. Currently, only Linux x86_64 is supported. + quiet || @error """No oneAPI driver found for your platform. Currently, only Linux x86_64 is supported. If you have a local oneAPI toolchain, you can use that; refer to the documentation for more details.""" return end @@ -211,7 +219,7 @@ function __init__() return end # For other errors, still report them as errors - @error "Failed to initialize oneAPI" exception=(err,catch_backtrace()) + quiet || @error "Failed to initialize oneAPI" exception=(err,catch_backtrace()) functional[] = false return end @@ -221,22 +229,22 @@ function __init__() drv_count = Ref{UInt32}(0) zeDriverGet(drv_count, C_NULL) if drv_count[] == 0 - @info "oneAPI initialized but no drivers found. oneAPI.jl will not be functional." + quiet || @info "oneAPI initialized but no drivers found. oneAPI.jl will not be functional." functional[] = false return end catch err - @error "Failed to enumerate oneAPI drivers" exception = (err, catch_backtrace()) + quiet || @error "Failed to enumerate oneAPI drivers" exception = (err, catch_backtrace()) functional[] = false return end - functional[] = true - validation_layer[] = parse(Bool, get(ENV, "ZE_ENABLE_VALIDATION_LAYER", "false")) parameter_validation[] = parse(Bool, get(ENV, "ZE_ENABLE_PARAMETER_VALIDATION", "false")) SCRATCH_HEDGE[] = parse_env_bool("ONEAPI_SCRATCH_HEDGE", true) - return sync_each_submission!(parse_env_bool("ONEAPI_SYNC_EACH_SUBMISSION", false)) + sync_each_submission!(parse_env_bool("ONEAPI_SYNC_EACH_SUBMISSION", false)) + functional[] = true + return end end diff --git a/src/compiler/precompile.jl b/src/compiler/precompile.jl index b2ece979f..61b1c6282 100644 --- a/src/compiler/precompile.jl +++ b/src/compiler/precompile.jl @@ -1,4 +1,4 @@ -using PrecompileTools: @compile_workload +using PrecompileTools: @setup_workload, @compile_workload # Warm up the GPUCompiler -> SPIR-V pipeline during precompilation so the first real # `zefunction` call is cheap. This doesn't need a GPU: SPIR-V codegen is device-independent @@ -51,3 +51,67 @@ if _precompile_lts ? SPIRV_LLVM_Translator_jll.is_available() : end end end + +# Exercise the runtime when a GPU can be used while precompiling: allocating, copying, +# launching and displaying all take seconds to compile on first use otherwise. +function precompile_runtime_kernel(a) + @inbounds a[1] = 1 + return +end +function precompile_runtime() + a = oneArray(Float32[1, 2, 3]) + Array(a) + show(IOContext(IOBuffer(), :limit => true), MIME"text/plain"(), a) + + # these compile kernels, which has the same issue on Julia < 1.12 as above + @static if VERSION >= v"1.12-" + @oneapi precompile_runtime_kernel(a) + Array(a .+ 1f0) + Array(oneArray([1]) .+ 1) + Array(rand(Float32, 2)) + end + synchronize() + return +end +# streams are task-local, so run the workload in a task of its own, whose storage can be +# collected along with the streams before the state they are registered in is reset. (from +# a function, as top-level code keeps the values it computes, like the task, alive) +function run_precompile_runtime() + wait(@async precompile_runtime()) + return +end + +@setup_workload begin + if can_precompile() + @compile_workload run_precompile_runtime() + + # the state the workload created refers to driver objects of this process. first + # collect the streams, whose finalizers need that state: free the workload's arrays, + # then drop the pools, whose freed chunks keep the streams that used them alive. + Base.@lock stream_registry_lock empty!(stream_registry) + Base.@lock pending_release_lock empty!(pending_release) + Base.@lock pending_streams_lock empty!(pending_streams) + GC.gc(true) + trim_pools!() + with_pool_lock(() -> empty!(pools)) + GC.gc(true) + + Base.@lock global_contexts_lock empty!(global_contexts) + Base.@lock batching_support_lock empty!(batching_support) + empty!(stagings) + empty!(staged_devices) + empty!(eltype_support_cache) + empty!(free_report_slots) + empty!(_compiler_configs) + empty!(oneL0.extension_properties_cache) + Base.@atomic oneAPIKernels.device_limits_cache.limits = oneAPIKernels.DeviceLimitsDict() + empty!(oneAPIKernels.pending_events) + _allocated_bytes[] = 0 + _total_mem_cache[] = 0 + _memory_limit_cache[] = -1 + end +end + +# whatever the probe found, keep it and the driver state out of the package image +precompile_probe[] = nothing +oneL0.functional[] = false diff --git a/src/oneAPI.jl b/src/oneAPI.jl index 299361546..46284ce8b 100644 --- a/src/oneAPI.jl +++ b/src/oneAPI.jl @@ -95,7 +95,46 @@ include("oneAPIKernels.jl") import .oneAPIKernels: oneAPIBackend export oneAPIBackend -# precompilation workload (warms up the SPIR-V compilation pipeline) +# whether precompilation workloads can use the GPU, also those of downstream packages like +# AcceleratedKernels' oneAPI extension. precompilation processes don't initialize the +# driver when loading oneAPI, so this does, once and without reporting why the GPU can't +# be used. disabled with the `precompile` preference. +const precompile_probe = Ref{Union{Nothing,Bool}}(nothing) +const precompile_probe_lock = ReentrantLock() +function can_precompile() + ccall(:jl_generating_output, Cint, ()) == 1 || return false + Preferences.@load_preference("precompile", true) || return false + @lock precompile_probe_lock begin + if precompile_probe[] === nothing + precompile_probe[] = probe_runtime() + end + return precompile_probe[]::Bool + end +end + +# whether the runtime can be used: the SPIR-V tool for the selected codegen path (see +# `_compiler_config`), and a device supporting the in-order immediate command lists that +# streams require (see `create_stream`) +function probe_runtime() + oneL0.initialize(; quiet=true) + functional() || return false + tool = oneL0.LTS[] ? SPIRV_LLVM_Translator_jll : SPIRV_LLVM_Backend_jll + tool.is_available() || return false + drv = first(drivers()) + devs = devices(drv) + isempty(devs) && return false + try + oneL0.ZeImmediateCommandList(ZeContext(drv), first(devs); + flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, + mode = oneL0.ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS) + catch err + err isa oneL0.ZeError || rethrow() + return false + end + return true +end + +# precompilation workloads (warming up the SPIR-V compilation pipeline, and the runtime) include("compiler/precompile.jl") # Work around a deadlock in Pkg's parallel precompilation on Julia 1.10, where it does diff --git a/test/level-zero.jl b/test/level-zero.jl index 0424801fa..958226903 100644 --- a/test/level-zero.jl +++ b/test/level-zero.jl @@ -7,6 +7,14 @@ if oneL0.NEO_jll.is_available() end +@testset "precompilation probe" begin + # what the precompilation workload found isn't kept in the package image, so it doesn't + # decide for other precompilation processes, and only applies while precompiling + @test oneAPI.precompile_probe[] === nothing + @test !oneAPI.can_precompile() +end + + @testset "driver" begin drvs = drivers()