Memoize the driver extension list and reduce launch-path overhead - #605
Open
michel2323 wants to merge 5 commits into
Open
Memoize the driver extension list and reduce launch-path overhead#605michel2323 wants to merge 5 commits into
michel2323 wants to merge 5 commits into
Conversation
Contributor
|
Your PR no longer requires formatting changes. Thank you for your contribution! |
michel2323
force-pushed
the
cache-launch-configuration
branch
from
August 7, 2026 13:38
7c4dcc1 to
5f74421
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #605 +/- ##
==========================================
+ Coverage 78.64% 80.22% +1.58%
==========================================
Files 50 50
Lines 3488 3490 +2
==========================================
+ Hits 2743 2800 +57
+ Misses 745 690 -55 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
`properties(::ZeKernel)` consults `extension_properties` to decide whether to link the max-group-size extension struct, and the KernelAbstractions backend reaches it through `launch_configuration` on every dispatch of a kernel with a dynamic workgroupsize. Rebuilding the answer each time takes two driver round-trips plus a `String` allocation and a `Dict` insertion per extension, so the query dominates the host cost of launch-bound workloads. A driver's extension list is fixed for its lifetime, and Level Zero has no driver-destroy entry point, so the result can be cached per driver. `ZeDriver` compares and hashes by handle, keying the cache by value, and the cache is emptied in `__init__` so no entry can be inherited from a precompiled image.
The `Base.fill!` specialization drove Level Zero's memory-fill command, which requires the fill 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. The repeated host allocations also pushed `zeMemAllocHost` into `retry_reclaim`'s garbage collections. GPUArrays' generic definition lowers to a single fill kernel and does none of this, so drop the specialization and let it apply.
`launch_configuration` returned the kernel's `maxGroupSize`, which Level Zero reports without regard to spill. The driver allocates `spillMemSize * group_size` bytes of scratch per work-group, so a heavily spilling kernel is reported as launchable at group sizes with a very large scratch demand. Cap the group size so a spilling kernel's per-group scratch stays within a fixed budget. Level Zero exposes no scratch-space query, so the budget is a conservative constant; CUDA.jl gets the equivalent from an occupancy API that accounts for register pressure. The inner division is clamped so a kernel spilling more than the whole budget still gets a group size of one.
`link_extensions` chains the extension descriptors together with raw interior pointers stored into each other's `pNext` fields, which the GC cannot see: `zeKernelGetProperties` receives only `props_ref` and reaches the other structs through the chain, so nothing in the call itself keeps them rooted. Preserve all three refs across the query, as `device_alloc`, `ZeModule` and `ZeKernel` already do for the buffers their descriptors point at. Defensive; not a fix for an observed failure.
michel2323
changed the title
michel2323
force-pushed
the
cache-launch-configuration
branch
2 times, most recently
from
August 7, 2026 15:40
25f9079 to
0690b28
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The KernelAbstractions backend reaches
properties(::ZeKernel)throughlaunch_configurationon every dispatch of a kernel with a dynamic workgroupsize, andextension_propertiesrebuilt the driver's extension list on each of those calls (the source carried aTODO: memoizethere). On a launch-bound workload that query dominates the per-launch host cost.Four independent changes:
__init__so nothing is inherited from a precompiled image).Base.fill!specialization required a USM host allocation, a residency call, a free and a full queue synchronize per call; GPUArrays' generic kernel fill does none of this.maxGroupSizeignoresspillMemSize, so a heavily spilling kernel was launched at group sizes with very large per-group scratch demands; cap the group size to a fixed scratch budget.pNextchain stores raw interior pointers the GC cannot see;GC.@preservethem like the neighbouring constructors do. Defensive, not a fix for an observed failure.(This PR originally cached
launch_configurationper kernel; memoizingextension_propertiessupersedes that — it removes the same cost and benefits every caller ofproperties.)