Repository navigation
POCL: run kernels on Julia's threads with PoCL host kernels - #906
Merged
Merged
Conversation
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 19, 2026 17:54
b3f62b1 to
93699f3
Compare
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 20, 2026 07:57
93699f3 to
f7a9ae0
Compare
Contributor
Benchmark ResultsShow table
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
Collaborator
|
This brings back some allocations that were removed with #901 👀 |
Member
Author
Should be fixed. |
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
2 times, most recently
from
October 20, 2026 13:42
a9c6833 to
dd53775
Compare
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 20, 2026 16:06
dd53775 to
7eef96c
Compare
vchuravy
reviewed
Oct 20, 2026
vchuravy
reviewed
Oct 20, 2026
Collaborator
|
I asked the bot to measure the impact of this PR on Oceananigans 👇 PR 906 cuts the cost of launching a kernel on the CPU by 10–50 times compared with KA 0.10.1, which brings it back to KA 0.9's level or better. Oceananigans workloads with many small kernels get 1.5–5 times faster. It has two costs, though: kernels that run for milliseconds get 7–33% slower with 4–8 threads, and each new kernel takes about 1.5 times as long to compile. All three rounds finished with no errors, using Julia 1.13 and PR head Launch cost, standalone kernels (µs per launch, at 1 / 4 / 8 threads)
Oceananigans (1 / 4 / 8 threads)
Memory allocations
Compile latency
The PR adds about 30–60 ms per new kernel on top of 0.10.1, which comes to about 4–9 s more before the first step of a real model finishes. Main is still far ahead here. Bottom line: the PR fixes the launch overhead KA 0.10 introduced and beats KA 0.9 on any kernel that isn't trivially small. Before it merges, two things need fixing upstream: the 7–33% slowdown of long kernels with several threads, and the extra compile latency. |
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 21, 2026 05:46
7eef96c to
9581115
Compare
Member
Author
|
Thanks. Analysis by Claude:
|
Collaborator
|
The machine I used for the benchmarks often has some bursts of high load, I'm running again the benchmarks now to see if the previous results still reproduce (load right now seems to be relatively low), anyway for the record below are the KA-only and Oceananigans benchmarks, the issue with longer kernels was allegedly visible already with the KA-only ones. Machine is Script to run the benchmarks is #!/bin/bash
# Each env on its own CCD (L3), CCDs rotated across rounds; output tagged with env/threads/round.
S=/path/to/scratchpad
OUT=$S/results.txt
ENVS=(ka09 ka010 ka906)
CCD=(16 64 72)
for r in 1 2 3; do
for i in 0 1 2; do
(
e=${ENVS[$i]}; base=${CCD[$(( (i + r) % 3 ))]}
for t in 1 4 8; do
cores="$base-$((base + t - 1))"
for script in kalat oclat; do
taskset -c $cores timeout 1800 julia +1.13 --startup-file=no -t $t --project=$S/env-$e $S/$script.jl > $S/logs/$script-$e-t$t-r$r.log 2>&1
grep -E '^RESULT|ERROR' $S/logs/$script-$e-t$t-r$r.log | sed "s/^/env=$e t=$t r=$r cores=$cores /" >> $OUT.$e
done
done
) &
done
wait
done
cat $OUT.ka09 $OUT.ka010 $OUT.ka906 > $OUT
echo done > $S/matrix.done
KernelAbstractions only benchmarks# Host cost, allocations and compile latency of KernelAbstractions kernels on the CPU backend.
# taskset -c <cores> julia +1.13 -t N --project=env kalat.jl
t_start = time()
using KernelAbstractions
using Statistics: median
using Printf
t_load = time() - t_start
@kernel function empty_kernel!(a)
end
@kernel function trivial!(a)
i = @index(Global, Linear)
@inbounds a[i] = 1f0
end
@kernel function many_args!(a, xs)
i = @index(Global, Linear)
@inbounds a[i] = first(xs)[i]
end
@kernel function stencil!(out, ins)
i, j, k = @index(Global, NTuple)
s = zero(eltype(out))
for x in ins
@inbounds s += x[i, j, k] + x[i + 1, j, k] - x[i, j + 1, k] + x[i, j, k + 1]
end
@inbounds out[i, j, k] = s / length(ins)
end
@kernel function new_kernel!(a, v, n)
i = @index(Global, Linear)
@inbounds a[i] = v
end
const backend = CPU()
sync() = KernelAbstractions.synchronize(backend)
# minimum and median over `reps` windows of the time per launch of `n` back-to-back launches
function per_launch(f, n; reps = 15)
ts = map(1:reps) do _
t0 = time_ns()
for _ in 1:n
f()
end
sync()
(time_ns() - t0) / n / 1e3
end
return minimum(ts), median(ts)
end
a1 = zeros(Float32, 1)
a66 = zeros(Float32, 66, 66)
small = zeros(Float32, 16, 16)
xs = ntuple(_ -> ones(Float32, 16, 16), 8)
out32 = zeros(Float32, 32, 32, 32)
ins32 = ntuple(_ -> rand(Float32, 33, 33, 33), 8)
out128 = zeros(Float32, 128, 128, 128)
ins128 = ntuple(_ -> rand(Float32, 129, 129, 129), 8)
cases = (
"empty, ndrange 1" => (() -> empty_kernel!(backend)(a1; ndrange = 1), 20_000),
"empty, ndrange 66²" => (() -> empty_kernel!(backend)(a66; ndrange = size(a66)), 20_000),
"trivial write, ndrange 1" => (() -> trivial!(backend)(a1; ndrange = 1), 20_000),
"trivial write, 66²" => (() -> trivial!(backend)(a66; ndrange = size(a66)), 20_000),
"8 arrays as args, 16²" => (() -> many_args!(backend)(small, xs; ndrange = size(small)), 20_000),
"stencil 8 arrays, 32³" => (() -> stencil!(backend)(out32, ins32; ndrange = size(out32)), 500),
"stencil 8 arrays, 128³" => (() -> stencil!(backend)(out128, ins128; ndrange = size(out128)), 10),
)
@printf("RESULT ka=%s threads=%d load_s=%.2f\n", pkgversion(KernelAbstractions), Threads.nthreads(), t_load)
for (name, (f, n)) in cases
t_first = @elapsed (f(); sync())
tmin, tmed = per_launch(f, n)
m = min(n, 2000)
f(); sync()
b = @allocated (for _ in 1:m; f(); end; sync())
@printf("RESULT case=%-26s first_ms=%9.1f min_us=%10.2f med_us=%10.2f B_per_launch=%8.1f\n",
repr(name), 1e3 * t_first, tmin, tmed, b / m)
end
# first launch (compilation included) of new kernels: each `Val(n)` is a new kernel
firsts = map(1:11) do n
@elapsed (new_kernel!(backend)(a66, 1f0, Val(n); ndrange = size(a66)); sync())
end
@printf("RESULT case=\"new kernel compile\" first_ms=%.1f median_next10_ms=%.1f\n", 1e3 * firsts[1], 1e3 * median(firsts[2:end]))Oceananigans benchmarks# Oceananigans workloads on the CPU: individual operations and full time steps.
# Reports first-call time (compilation), steady time per call and bytes allocated per call.
# taskset -c <cores> julia +1.13 -t N --project=env oclat.jl
t_start = time()
using Oceananigans
using Oceananigans.BoundaryConditions: fill_halo_regions!
using Statistics: median
using Printf
using KernelAbstractions
t_load = time() - t_start
@printf("RESULT ka=%s threads=%d load_s=%.2f\n", pkgversion(KernelAbstractions), Threads.nthreads(), t_load)
# `windows` windows of `n` calls; minimum and median time per call
function steady(f, n; windows = 10)
for _ in 1:3; f(); end
ts = map(1:windows) do _
(@elapsed for _ in 1:n; f(); end) / n
end
b = @allocated f()
return minimum(ts), median(ts), b
end
function report(name, setup_s, f, n)
t_first = @elapsed f()
tmin, tmed, b = steady(f, n)
unit, s = tmin < 1e-3 ? ("us", 1e6) : ("ms", 1e3)
@printf("RESULT case=%-34s setup_s=%7.2f first_s=%7.2f min_%s=%10.3f med_%s=%10.3f KiB_per_call=%8.2f\n",
repr(name), setup_s, t_first, unit, s * tmin, unit, s * tmed, b / 1024)
end
for N in (16, 128)
t_setup = @elapsed begin
grid = RectilinearGrid(CPU(); size = (N, N, N), extent = (1, 1, 1), halo = (3, 3, 3))
c = CenterField(grid)
set!(c, (x, y, z) -> x * y + z)
op = Field(∂x(c) * ∂y(c) + c^2)
end
n = N == 16 ? 2000 : 20
report("fill_halo_regions! $(N)³", t_setup, () -> fill_halo_regions!(c), n)
report("compute! ∂x(c)*∂y(c)+c² $(N)³", 0.0, () -> compute!(op), n)
end
function hydrostatic(Nx, Ny, Nz)
grid = LatitudeLongitudeGrid(CPU(); size = (Nx, Ny, Nz), longitude = (0, 360), latitude = (-75, 75),
z = (-1000, 0), halo = (7, 7, 7))
model = HydrostaticFreeSurfaceModel(grid;
momentum_advection = WENOVectorInvariant(),
tracer_advection = WENO(order = 7),
buoyancy = SeawaterBuoyancy(),
tracers = (:T, :S),
closure = CATKEVerticalDiffusivity(),
coriolis = HydrostaticSphericalCoriolis(),
free_surface = SplitExplicitFreeSurface(grid; substeps = 30))
set!(model; T = (λ, φ, z) -> 20 + 10 * z / 1000 + 0.01 * rand(), S = 35)
return model
end
function nonhydrostatic(N)
grid = RectilinearGrid(CPU(); size = (N, N, N), extent = (1, 1, 1), halo = (5, 5, 5))
model = NonhydrostaticModel(grid; advection = WENO(order = 9), buoyancy = BuoyancyTracer(), tracers = :b,
closure = ScalarDiffusivity(ν = 1e-4, κ = 1e-4))
set!(model; b = (x, y, z) -> z + 1e-3 * rand(), u = (x, y, z) -> 1e-3 * rand())
return model
end
for (name, build, n) in (("hydrostatic latlon 36x18x10", () -> hydrostatic(36, 18, 10), 20),
("hydrostatic latlon 180x90x10", () -> hydrostatic(180, 90, 10), 5),
("nonhydrostatic 32³", () -> nonhydrostatic(32), 20))
t_setup = @elapsed (model = build())
report("time_step! " * name, t_setup, () -> time_step!(model, 1), n)
end |
Collaborator
|
Results on new head of the branch. I'll ask it to investigate the slowdown for the larger hydrostatic model at 4-8 threads, but it'd be great if you could at least try to reproduce it. I'm using CliMA/Oceananigans.jl#5799 (of course you'll have to point KA to this branch). The update cut the PR's extra compile time per kernel by about half and left everything else unchanged. Launch cost, allocations and almost all runtimes are the same as with the old head. The new head also looks steadier on the small hydrostatic model. The 4–8 thread slowdown of the larger hydrostatic model compared with KA 0.10.1 is still there. All four environments ran together in each of 3 rounds, so they share today's host load. Columns are KA 0.9 (main) / KA 0.10.1 / old PR head Launch cost, standalone kernels (µs per launch, at 1 / 4 / 8 threads)
Oceananigans (1 / 4 / 8 threads)
* KA 0.9 was noisy on this row today (4.2–13.3 ms at 4 threads), which looks like host interference.
Memory allocationsThese are identical for the old head, the new head and KA 0.10.1:
Main with KA 0.9 still allocates more as threads increase: 2.3–5 KB per launch and 0.5–1 MiB per step at 4–8 threads. Compile latency
Bottom line: the update made compilation cheaper and didn't change launch cost. Before this is better than 0.10.1 everywhere, two things remain:
|
Member
Author
|
Thanks. Analysis by Claude:
|
Collaborator
|
Do I understand correctly that this now uses the interactive thread when available, so that if we do thread pinning we either need to use |
Collaborator
|
New summary, I like that it put the "bottom line" at the top.
Edit: it turned out that the |
Member
Author
Partly, yes. Launches are sized to
So: |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #906 +/- ##
==========================================
- Coverage 84.95% 84.62% -0.33%
==========================================
Files 31 31
Lines 2612 2901 +289
==========================================
+ Hits 2219 2455 +236
- Misses 393 446 +53 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 22, 2026 11:57
9581115 to
7e60b1b
Compare
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 23, 2026 13:13
7e60b1b to
b1cf2ca
Compare
14 tasks
This was referenced Oct 27, 2026
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 28, 2026 08:00
ddd6405 to
6290988
Compare
This was referenced Oct 28, 2026
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
October 28, 2026 09:15
6290988 to
d24bda8
Compare
Contributor
Benchmark resultsCompared with 8a985320e, the merge base. Times are the minimum time of every benchmark. A change is unusual when it exceeds both three times the benchmark's typical change between consecutive runs on the default branch and a minimum per category (performance 3%, latency 5%, applications 3%, other 3%). A benchmark that starts or stops allocating is listed however few the allocations. A single run can still be a fluke, so rerun before drawing conclusions.
All benchmarks
|
vchuravy
approved these changes
Oct 30, 2026
vchuravy
left a comment
vchuravy
left a comment
Member
There was a problem hiding this comment.
Generally looks good, but didnt have time to look at it in detail
vchuravy
reviewed
Oct 30, 2026
| return launch_kernel(obj, launch, subgroups, ndrange, _workgroupsize, iterspace, args, true) | ||
| extent = split_extent(size(cartesian(range)), items) | ||
| if extent !== nothing && split_launch(b, launch, kernel) | ||
| full_tt = Base.compilerbarrier(:type, without_bounds_check(tt))::Type |
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
3 times, most recently
from
November 6, 2026 08:00
98ca588 to
5986a9e
Compare
Handing a launch to PoCL's thread pool and waiting for it costs about 5 µs, even for an empty kernel. PoCL's `cl_pocl_host_kernel` extension, in every pocl_standalone_jll since 7.2.3, runs work-groups on the calling thread instead. On a 4-core x86 machine, a launch now costs about 0.15 µs on the launching task and 1 µs on a ThreadTeams team. Launches of up to 8192 work-items stay on the launching task unless the kernel's recent launches were slow. Larger ones use `foreach_range`, which starts block `k` of a `Threads.@threads :static` split on thread `k`, keeping data in place across launches; threads steal work when the grid changed or the previous launch took 50 µs or more. ThreadTeams 1.3 steals from nearby threads first. Host kernels can't yield: interrupts are deferred until the launch ends, and tasks, timers and I/O run after team and long launches. Runs are GC-safe, and launches from finalizers, which can't switch tasks, run every work-group on the calling thread. The command queue remains for explicit work-group sizes over the 1 MiB stack limit, kernels PoCL can't prepare, and direct POCL calls with non-uniform sizes. The thread count is capped at Julia's default pool.
Only partial work-groups need the bounds check, but a launch with any ran the checked kernel, whose work-item loop vectorizes poorly, for all of them. Run the full ones with a variant compiled without the check, in the same launch, sharing arguments, geometry and exception status. Linear launches, launches on the command queue, and launches past three dimensions with a partial third dimension stay checked throughout. A kernel's first launch with partial work-groups stays checked, so that kernels used once, such as AcceleratedKernels closures, compile once: on 4 Xeon E3-1245 v6 threads, a trivial kernel's first launch takes 122 ms, as through PoCL's queue, instead of 167 ms. If the unchecked variant can't run as a host kernel, the checked one runs every work-group.
Work-groups whose stack estimate exceeds 1 MiB run on PoCL's command queue instead of as host kernels. WaterLily's donut body kernel (344 B per work-item, 4096 work-items on Zen 5) did so every step. Cap tuned work-groups, once per kernel at link time, at the largest power of two that fits; explicit sizes still fall back. On 16 Zen 5 threads, donut_5 takes 2.2-2.35 ms per step instead of 2.9-3.0, and donut_6 11.3-13.2 ms instead of 14.4-15.0. Fitting work-groups to the ndrange also left small work-item loops with scalar remainders: axpy over 1000 work-items on 32 threads took 162 ns with groups of 25 instead of 146 ns with groups of 32. Shrink a first dimension below 16 vectors only to a vector-width multiple.
Describe which thread runs which work-groups, when a launch stays on the launching task or lets other tasks run, and which kernels still run on PoCL's threads. Pin threads with ThreadPinning.jl rather than `POCL_AFFINITY`; the NUMA example and its measurements now do so.
A launch defers `^C` until all work-groups finish, but ThreadTeams' dispatcher yields while it waits for the team. On Julia 1.13 and earlier, a task switch delivers the deferred SIGINT to whichever task runs next on thread 1, usually ThreadTeams' worker there, so the launch returned normally and the interrupt was lost. Dispatch with ThreadTeams 1.4's `isolate = true`, which never switches tasks, on all versions. The first scratch slot lookup on a thread may also take a lock and switch tasks; create the slots on the launching task before a team launch starts. Document that before Julia 1.14 only launches from thread 1 are protected, and test a launch with a task waiting on its thread.
The small host launches test expected one of five team launches of a trivial kernel to estimate a serial duration under 4 µs. A team run times its first range on a team thread, which on CI may have just woken up or been descheduled: the macOS runners have 3 cores for 5 Julia threads. One CI run measured five slow ranges in a row, so the kernel stayed on the team. Launch until the schedule turns serial, up to 100 times; the first fast run is enough.
maleadt
force-pushed
the
tb/pocl-host-kernel-bench
branch
from
November 7, 2026 11:42
e08f7dd to
d836f57
Compare
On macOS, Julia before 1.13 installs SIGINT's handler on the main thread while the signal listener thread sets it to SIG_IGN for kqueue. When the listener loses that race, the process ignores every SIGINT, and all six interrupt testsets fail with `err === nothing` (seen on #941's macOS 1.10 CI). Fixed upstream in JuliaLang/julia#62471, which only 1.13 has.
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.
The
CPUback-end runs work-groups on Julia's threads through PoCL'scl_pocl_host_kernelextension (pocl_standalone_jll7.2.3 or later). Launches bypass PoCL's command queue, events and thread pool, so launching a small kernel costs 0.1–0.3 µs instead of 4–82 µs with KA 0.10.3. Larger kernels run about as fast as with Polyester, or faster (see Known gaps for exceptions).How it works
ThreadTeams.foreach_range, a team of Julia tasks that stays warm between launches.Threads.@threads :staticsplit: blockkstarts on default-pool threadk, regardless of the launching thread. A launch keeps this static split if the previous launch took < 50 µs and had the same group count. Otherwise, including on the first launch, idle threads can steal work so a slow or descheduled thread doesn't hold up the launch. ThreadTeams 1.3 steals from nearby threads first.KernelAbstractions.zeros/onesand the NUMA example therefore initialize data with a kernel.ndrange, or otherwise to a quarter, preserving vector-width multiples in small first dimensions. They also shrink to fit a host kernel's 1 MiB stack allowance on a Julia task.isolate = true, which never switches tasks. Before Julia 1.14,^Cgoes to whatever runs on thread 1, so only launches from there (e.g. the main task) keep it. Other tasks, timers and I/O run after launches that use the team or run long. Launches run on the launching task alone when called from a finalizer or aThreads.@threadsloop, or when another task's launch is using the team.PoCL's threads still run kernels whose explicit work-group sizes require more than the 1 MiB stack limit, kernels PoCL can't prepare as host kernels, and direct POCL calls with non-uniform sizes.
JULIA_KA_CPU_THREADScan lower the thread count but cannot exceed Julia's default-pool size.Pinning
Pin Julia's threads before the first launch: ThreadTeams reads the topology when its threads start. Team launches run block
kon default-pool threadk. A launch from the interactive thread (where Julia 1.12+ runs the main task by default) runs no block itself unless a worker starts late, but small launches run entirely on it. So pin the launching thread too: start Julia with-t N,0and use ThreadPinning.jl'spinthreads(:cores)orpinthreads(:numa), or keep the interactive thread and passthreadpool = :alltopinthreads. With fewer threads than cores,pinthreads(:numa)spreads threads over all NUMA domains.POCL_AFFINITY=1affects only kernels running on PoCL's threads.tasksetalone isn't enough: on a Ryzen 5950X withtaskset -c 0-15and-t 16, Linux sometimes packs all 16 threads onto one CCD, and axpy 1e5 takes 5–9 µs instead of 3.Results
Median µs per launch at d24bda8's predecessor 6290988 (same performance within noise), using Julia 1.12.7 and ThreadTeams 1.3.0. Cells show this PR (best Polyester / KA 0.10.3).
Workstations,
-t N, threads restricted withtaskseton Linux:* Bimodal: 2.9–3.3 µs per process unless the scheduler packs the threads onto one CCD (see Pinning).
Servers,
-t N,0,pinthreads(:cores):At 32 threads,
:coresuses one socket of the two-socket machines.NUMA, ThreadTeams 1.3 against 1.2.1 at 64 threads, axpy 1e7, µs:
:cores:numaSAXPY over 2 × 500M
Float32values initialized by a kernel: 242 GB/s at 64 threads on the EPYC (both versions;Threads.@threads :staticgets 240). At 32 threads,:numaspreads the threads over both sockets: 157 GB/s instead of 122 with:coreson the EPYC, 343–360 instead of 216 on the Xeon.With a competing process on one of 4 cores (Xeon E3,
-t 4,0), axpy 1e6 takes 79.7 µs (Polyester 108, KA 0.10.3 202).Applications, ms per time step at an earlier head (f9d9337, ThreadTeams 1.2.0). Cells show this PR (KA 0.10.3; Polyester port of WaterLily's
@loop):Since those measurements, fitting work-groups to the host-kernel stack has kept WaterLily's donut body kernel off PoCL's threads: 11.3–13.2 ms per step instead of 14.4–15.0 on 16 Zen 5 threads.
Known gaps
Related
main.foreach_range(...; schedule = :steal): prefer victims sharing a cache or NUMA node vchuravy/ThreadTeams.jl#85 (ThreadTeams 1.3): nearby-first stealing.foreach_range(...; isolate = true): run no other task while waiting vchuravy/ThreadTeams.jl#87 (ThreadTeams 1.4):foreach_range(...; isolate = true), so that waiting for the team doesn't hand a deferred^Cto another task.