Memoize the driver extension list and reduce launch-path overhead

lib/level-zero/driver.jl

@@ -100,7 +100,34 @@ function ipc_properties(drv::ZeDriver)
     )
 end
 
+# Memoized, because this lands on the per-launch path: `properties(::ZeKernel)` consults it
+# to decide whether to link the max-group-size extension struct, and KernelAbstractions
+# reaches that from `launch_configuration` on every dispatch of a kernel whose
+# workgroupsize is dynamic. Answering it is not cheap -- two driver round-trips, a copy of
+# the property vector, and a `String` allocation plus a `Dict` insertion per extension --
+# and on a Data Center GPU Max 1550, whose driver advertises many extensions, it measured
+# ~90 us per call and accounted for 90% of the host time of a launch-bound GPU callback.
+#
+# Safe to remember: a driver reports a fixed set of extensions, and Level Zero has no
+# driver-destroy entry point, so a `ZeDriver` handle stays valid and keeps its meaning for
+# the lifetime of the process. The cache is emptied in `__init__` so that entries can never
+# be inherited from the process that generated a precompiled image.
+#
+# The returned `Dict` is the cached object itself; callers must treat it as read-only.
+const extension_properties_cache = Dict{ZeDriver, Dict{String, VersionNumber}}()
+const extension_properties_lock = ReentrantLock()
+
 function extension_properties(drv::ZeDriver)
+    Base.@lock extension_properties_lock begin
+        cached = get(extension_properties_cache, drv, nothing)
+        cached === nothing || return cached
+        extensions = _extension_properties(drv)
+        extension_properties_cache[drv] = extensions
+        return extensions
+    end
+end
+
+function _extension_properties(drv::ZeDriver)
     count_ref = Ref{UInt32}(0)
     zeDriverGetExtensionProperties(drv, count_ref, C_NULL)
 

lib/level-zero/module.jl

@@ -235,16 +235,27 @@ export properties
 function properties(kernel::ZeKernel)
     props_ref = Ref(ze_kernel_properties_t())
     preferred_group_size_props_ref = Ref(ze_kernel_preferred_group_size_properties_t())
-    link_extensions(props_ref, preferred_group_size_props_ref)
-    if haskey(oneL0.extension_properties(kernel.mod.context.driver),
-              "ZE_extension_kernel_max_group_size_properties")
-        # TODO: memoize
-        max_group_size_props_ref = Ref(ze_kernel_max_group_size_properties_ext_t())
-        link_extensions(preferred_group_size_props_ref, max_group_size_props_ref)
+    max_group_size_props_ref =
+    if haskey(
+            oneL0.extension_properties(kernel.mod.context.driver),
+            "ZE_extension_kernel_max_group_size_properties"
+        )
+        Ref(ze_kernel_max_group_size_properties_ext_t())
     else
-        max_group_size_props_ref = nothing
+        nothing
+    end
+
+    # `link_extensions` chains these together with raw interior pointers stored into each
+    # other's `pNext` field, which the GC cannot see: only `props_ref` is passed to the
+    # query, and the driver reaches the rest by following those pointers. Hold them across
+    # the call, as `device_alloc` does for its relaxed-allocation extension.
+    GC.@preserve props_ref preferred_group_size_props_ref max_group_size_props_ref begin
+        link_extensions(props_ref, preferred_group_size_props_ref)
+        if max_group_size_props_ref !== nothing
+            link_extensions(preferred_group_size_props_ref, max_group_size_props_ref)
+        end
+        zeKernelGetProperties(kernel, props_ref)
     end
-    zeKernelGetProperties(kernel, props_ref)
 
     props = props_ref[]
     return (

lib/level-zero/oneL0.jl

@@ -166,6 +166,10 @@ function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
     precompiling && return
 
+    # Driver handles are process-local, so nothing keyed on one may survive into another
+    # process through a precompiled image.
+    empty!(extension_properties_cache)
+
     # Resolve the LTS master switch up front, before the driver-availability early
     # returns below: it gates codegen and behavior and must be set even on hosts
     # without a functional GPU. Default off (rolling stack); an LTS deployment such as

src/array.jl

@@ -544,23 +544,11 @@ ones(dims...) = ones(Float64, dims...)
 fill(v, dims...) = fill!(oneArray{typeof(v)}(undef, dims...), v)
 fill(v, dims::Dims) = fill!(oneArray{typeof(v)}(undef, dims...), v)
 
-function Base.fill!(A::oneDenseArray{T}, val) where T
-  length(A) == 0 && return A
-  val = convert(T, val)
-  sizeof(T) == 0 && return A
-
-  # execute! is async, so we need to allocate the pattern in USM memory and keep it alive
-  # until the operation completes. The fill reads this host buffer on the GPU, so it must
-  # be made resident on the device like any other USM a kernel reads (see
-  # `allocate(::Type{oneL0.HostBuffer}, ...)`).
-  buf = oneL0.host_alloc(context(A), sizeof(T), Base.datatype_alignment(T))
-  oneL0.make_resident(context(A), device(), buf)
-  unsafe_store!(convert(Ptr{T}, buf), val)
-  unsafe_fill!(context(A), device(), pointer(A), convert(ZePtr{T}, buf), length(A))
-  synchronize(global_queue(context(A), device()))
-  oneL0.free(buf)
-  A
-end
+# NOTE: `Base.fill!` is deliberately not specialized here. GPUArrays' generic definition
+# lowers to a single fill kernel, whereas the Level Zero memory-fill command requires the
+# pattern to live in USM memory: a host allocation, a residency call and a free around
+# every call, plus a full queue synchronize to keep the pattern alive until the
+# asynchronous fill has read it.
 
 
 ## derived arrays

src/compiler/execution.jl

@@ -209,6 +209,18 @@ struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     fun::ZeKernel
 end
 
+# Upper bound on the spill (scratch) memory a single work-group may require, in bytes.
+#
+# The driver allocates `spillMemSize * group_size` of scratch per work-group. `maxGroupSize`
+# does not account for spill, so a heavily spilling kernel is reported as launchable at a
+# group size whose scratch demand is enormous: on a Data Center GPU Max 1550, a kernel
+# spilling 3648 B/thread is reported launchable at 1024 items/group, i.e. ~3.7 MB of scratch
+# for one work-group.
+#
+# Level Zero exposes no scratch-space query, so the budget is a conservative constant rather
+# than a derived one.
+const MAX_GROUP_SCRATCH = 1024 * 1024
+
 function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
     # Level Zero's zeKernelSuggestGroupSize provides a launch configuration
     # that exactly cover the input size. This can result in very awkward
@@ -230,6 +242,16 @@ function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
         group_size = max_size ÷ 2
     end
 
+    # keep a spilling kernel's per-group scratch within budget. CUDA.jl gets this for free
+    # from an occupancy API that is register-pressure aware; Level Zero reports the spill
+    # size but does not fold it into `maxGroupSize`, so account for it here. Rounded down to
+    # a power of two, both because group sizes want to be anyway and to stay clear of the
+    # limit rather than right at it.
+    spill = kernel_props.spillMemSize
+    if spill > 0 && group_size * spill > MAX_GROUP_SCRATCH
+        group_size = max(1, prevpow(2, max(1, MAX_GROUP_SCRATCH ÷ spill)))
+    end
+
     # TODO: align the group size based on preferredGroupSize
 
     return group_size

test/level-zero.jl

@@ -22,6 +22,11 @@ properties(drv)
 ipc_properties(drv)
 extension_properties(drv)
 
+    # the extension list is memoized: it is answered from the cache, and that answer is the
+    # one the driver gives
+    @test extension_properties(drv) === extension_properties(drv)
+    @test extension_properties(drv) == oneL0._extension_properties(drv)
+
 end
 
 drv = first(drivers())