From 7c4dcc1a06571d9fd7beda47745395bd1c1871bc Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Wed, 5 Aug 2026 14:20:01 +0000 Subject: [PATCH] Cache launch_configuration per kernel `launch_configuration` issues a `zeKernelGetProperties` round-trip, and KernelAbstractions calls it on every dispatch of a kernel whose workgroupsize is `DynamicSize` and unspecified at the call site (src/oneAPIKernels.jl). The query therefore lands on the per-launch path, where it dominates: on a Data Center GPU Max 1550, a no-op kernel costs 96 us per launch as-is and 9.7 us with the workgroupsize pinned, a 9.9x difference that is entirely host-side. Launch-bound workloads pay it on every kernel they dispatch. The returned group size depends only on the kernel and its device. A ZeKernel's handle and module are fixed for its lifetime, and maxGroupSize / maxTotalGroupSize are static properties of the compiled kernel, so neither the ndrange nor any argument value can change the result. Keyed weakly on the ZeKernel rather than on its raw handle: a handle is unique only among live kernels, so a handle-keyed entry could be inherited by a later kernel that reuses a destroyed kernel's address. A weak key ties the entry to the kernel's lifetime and needs no hook in oneL0's finalizer. Measured with the cache in place: 96 us -> 62 us per launch. The remainder is the re-partition and context rebuild that follow the query in the KernelAbstractions backend, which this does not address. --- src/compiler/execution.jl | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) 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