Precompile the runtime when a GPU can be used while precompiling

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

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

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

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

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()