Repository navigation
Reduce the host code compiled for every new kernel - #565
Merged
Merged
Conversation
`kernel_context`, `kernel_device`, `throw_device_mismatch` and the body of `KI.max_work_group_size` only use the kernel's `cl.Kernel` handle, but took the `KI.Kernel`, whose type includes the kernel function and argument types, so they were inferred and compiled again for every kernel. Type them on the `cl.Kernel` instead. `KI.launch` and `KI.kernel_function` are necessarily compiled per kernel: move the activation and context check of the former into `prepare_launch`, and don't inline `activate` in the latter. Per new broadcast kernel, after warm-up: compile time 176 -> 143 ms, inference 178 -> 145 ms.
`cl.call` is compiled for every kernel signature, and inlined the loop over the indirect memory objects (union-split over the memory types) and the `clSetKernelExecInfo` calls, none of which depend on the arguments. Move them into a function that's compiled once, as oneAPI.jl does for its argument-independent launch code. Per new broadcast kernel, after warm-up: compile time 143 -> 116 ms, inference 145 -> 121 ms.
`launch_with_exception_mailbox` is compiled for every kernel signature, and inlined the mailbox lookup (a `get!` with a closure, under a lock), the launch ID atomics and the bookkeeping of pending queues. None of that depends on the arguments, so move it into functions that are compiled once. The calls, their order and the locks held around them are unchanged. Per new broadcast kernel, after warm-up: compile time 116 -> 106 ms, inference 121 -> 108 ms.
`clfunction` is specialized on the kernel function and argument types, but only the method instance and the returned `HostKernel` depend on them. Move the rest, i.e. the compiler configuration, the cache lookup and the per-context kernel scan and linking, into a function that takes the method instance, as oneAPI.jl does with `zefunction_link`. The method instance is now looked up before taking the compiler lock rather than under it. Per new broadcast kernel, after warm-up: compile time 106 -> 90 ms, inference 108 -> 96 ms.
`set_arg!` is compiled for every argument type, and contained the `clSetKernelArg` call and its error handling, including the interpolated message about mismatched argument sizes. Only compute the size and address of the argument there, and leave the rest to a function that's compiled once, as oneAPI.jl does for `zeKernelSetArgumentValue`. Per new broadcast kernel, after warm-up: compile time 90 -> 84 ms, inference 96 -> 86 ms.
KernelAbstractions' `launch_kernel` is compiled for every kernel, and inlined OpenCL.jl's `KI.max_work_group_size(::KI.Kernel)` (twice, through `launch_configuration` and `check_work_group_size`) as well as the task-local cache lookup of `device_properties`. Make the former forward to a function on the `cl.Kernel`, and don't inline either. Per new broadcast kernel, after warm-up: compile time 84 -> 78 ms, inference 86 -> 77 ms.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #565 +/- ##
==========================================
+ Coverage 87.52% 87.63% +0.10%
==========================================
Files 18 18
Lines 1716 1730 +14
==========================================
+ Hits 1502 1516 +14
Misses 214 214 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This was referenced Oct 16, 2026
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
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.
Since the move to KernelAbstractions 0.10 and KernelInterface (#519), the test suite takes about 30% longer in CI (~10.5-11 ks to ~14.5-14.8 ks summed). Apart from the new and larger KA test suites, the GPUArrays tests account for ~2000 s of that, at +10-50% each. That time goes to compiling host code: every new kernel signature recompiled the launch path, because the functions specialized on the
KI.Kernel/HostKerneltype or on the argument tuple inlined a lot of code that doesn't depend on either. For one new broadcast kernel, after warm-up, host compilation took 176 ms on main.This PR adds function barriers so that the argument-independent work is compiled once:
kernel_context,kernel_device,throw_device_mismatch,max_work_group_size) take thecl.Kernelhandle instead of theKI.Kernel;KI.launchleaves activation and its context check to a non-inlinedprepare_launch, andKI.kernel_functiondoesn't inlineactivate;cl.callregisters the indirect memory (the USM/SVM pointer loop andclSetKernelExecInfo) in a separate function;launch_with_exception_mailbox;clfunctiononly looks up the method instance and builds theHostKernel; the configuration, cache lookup and linking happen inclfunction_link(::MethodInstance);set_arg!only computes the size and address of the argument;clSetKernelArgand its error handling are in an unspecialized setter;This is the structure oneAPI.jl already uses (
launch_kernel(::ZeKernel, ...),set_argument_value!,zefunction_link). The calls happen in the same order as before, with one exception: the method instance is now looked up before taking the compiler lock, as in oneAPI.jl. Metal.jl instead puts@nospecializeon its launch body and dispatches dynamically to argument encoding, which would add a dynamic call to every launch here, so I didn't follow that.Compile time for a new broadcast kernel, after warm-up (Julia 1.12, PoCL 7.2, median of 3 interleaved runs):
Each commit message has its own measurement. On the
broadcastingtestset, a cold run takes 187-196 s on main and 130-137 s with this PR.What's left per kernel is the argument conversion, the keyword-call layers between
KI.launchandcl.call, and KA 0.10's ownlaunch_tuple/launch_kernel/prepare_arguments(~18 ms of inference). The latter is shared with all back-ends and needs a KernelAbstractions change; a prototype that moves its argument-independent setup and geometry into@noinlinehelpers takes another ~10 ms of compile time off.Tested locally on PoCL 7.2 with Julia 1.12 and 1.13 (
pocl/execution, kernel, exceptions, device, device/random, kernelinterface, kernelabstractions and gpuarrays/broadcasting, pluspoclc/execution and exceptions).