diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index cc8d3310..c3a74299 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -209,7 +209,37 @@ struct HostKernel{F,TT} <: AbstractKernel{F,TT} fun::ZeKernel end +# Memoized: this sits on the per-launch path, not on a setup path. +# +# KernelAbstractions calls `launch_configuration` on every dispatch of a kernel +# whose workgroupsize is `DynamicSize` and unspecified at the call site (see +# src/oneAPIKernels.jl), so the `zeKernelGetProperties` round-trip below is paid +# per launch. It dominates the launch: measured on a Data Center GPU Max 1550 +# with a no-op kernel, 96 us per launch with the query versus 9.7 us with the +# workgroupsize pinned, all of it host-side. Launch-bound workloads pay this on +# every kernel they dispatch. +# +# Safe to cache: the result is a function of the kernel and its device alone. +# A ZeKernel's handle and module are fixed for its lifetime, and `maxGroupSize` / +# `maxTotalGroupSize` are static properties of the compiled kernel -- neither the +# ndrange nor any argument value reaches this computation. +# +# Keyed weakly on the ZeKernel rather than on its raw handle: a handle is only +# unique among *live* kernels, so a handle key could let a destroyed kernel's +# entry be inherited by a later kernel that reuses its address. The weak key ties +# the entry to the kernel's own lifetime and needs no hook in oneL0's finalizer. +const _launch_config_cache = WeakKeyDict{ZeKernel,Int}() + function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT} + fun = kernel.fun + cached = get(_launch_config_cache, fun, nothing) + cached === nothing || return cached + config = _launch_configuration_uncached(kernel) + _launch_config_cache[fun] = config + return config +end + +function _launch_configuration_uncached(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 # configurations, so roll our own version that behaves like CUDA's