diff --git a/docs/src/installation.md b/docs/src/installation.md index 4f60a13a..5da810db 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 565391f3..d654abd7 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 b2ece979..61b1c628 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 29936154..46284ce8 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 0424801f..95822690 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()