Skip to content

Reduce the host code compiled for every new kernel - #565

Merged
maleadt merged 6 commits into
mainfrom
tb/launch-compile-time
Oct 16, 2026
Merged

maleadt merged 6 commits into
mainfrom
tb/launch-compile-time

Conversation

@maleadt

@maleadt maleadt commented Oct 16, 2026

Copy link
Copy Markdown
Member

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/HostKernel type 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:

  • the KernelInterface helpers (kernel_context, kernel_device, throw_device_mismatch, max_work_group_size) take the cl.Kernel handle instead of the KI.Kernel; KI.launch leaves activation and its context check to a non-inlined prepare_launch, and KI.kernel_function doesn't inline activate;
  • cl.call registers the indirect memory (the USM/SVM pointer loop and clSetKernelExecInfo) in a separate function;
  • the exception mailbox lookup, launch ID and pending-queue bookkeeping move out of launch_with_exception_mailbox;
  • clfunction only looks up the method instance and builds the HostKernel; the configuration, cache lookup and linking happen in clfunction_link(::MethodInstance);
  • set_arg! only computes the size and address of the argument; clSetKernelArg and its error handling are in an unspecialized setter;
  • KA's launch code calls non-inlined helpers for the kernel's work-group limit and the cached device properties.

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 @nospecialize on 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):

compile inference
main 176 ms 178 ms
this PR 78 ms 77 ms

Each commit message has its own measurement. On the broadcasting testset, 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.launch and cl.call, and KA 0.10's own launch_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 @noinline helpers 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, plus poclc/ execution and exceptions).

`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

codecov Bot commented Oct 16, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87.63%. Comparing base (d2b2f1f) to head (c7a88d0).

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant