[giordano]

Stacked on #941 (and #906).

A CPU launch without a workgroup size currently follows PoCL's pocl_default_local_size_optimizer, which stops splitting once there are more than half as many work-groups as threads. With #906's scheduling (block k on thread k), that leaves threads idle whenever the few large work-groups don't divide evenly among them. For example, a 132×132 launch on 8 threads ends up as 6 work-groups of up to 4096 work-items: two threads have nothing to do, and a thread that finishes early can't pick up more. 2-D launches whose work-items loop over a column, like the vertical integrals in ocean models, hit this often.

This PR keeps halving the work-group while there are at most four work-groups per thread. The floor stays at 4× the preferred work-group size multiple, so work-groups stay large enough to vectorize. Launches that already have enough work-groups are unchanged. 1-thread launches get smaller work-groups but run the same work serially.

Performance

A 2-D launch whose work-items sum and update a column, on a Threadripper PRO 7995WX pinned to one CCD (julia -t N,0):

launch	8 threads, #941	8 threads, this PR	16 threads, #941	16 threads, this PR
132×132×32	81–82 µs	67 µs (−18%)	48–49 µs	38 µs (−22%)
100×100×50	98 µs	72 µs (−27%)	51–52 µs	40 µs (−22%)
256×256×16	1.06 ms	1.05–1.06 ms	0.54 ms	0.51–0.55 ms

With 1 thread, all three are unchanged.

Oceananigans uses its own row-spanning CPU work-group to avoid this problem. With this PR, letting KernelAbstractions choose gives hydrostatic and nonhydrostatic time steps (32³ to 128²×32, 1 to 16 threads) within ±2% of that custom choice, and up to 11% faster for small grids at 8–16 threads. With #941's rule they were up to 23% slower.

Without #906's scheduling (KernelAbstractions 0.10.5), more work-groups per thread made launches slower at 16 threads, so this change belongs on top of #906.

Reproducer

using KernelAbstractions
using Printf

@kernel function column_sum!(out, @Const(a))
    i, j = @index(Global, NTuple)
    s = zero(eltype(out))
    for k in axes(a, 3)
        @inbounds s += a[i, j, k]
    end
    for k in axes(a, 3)
        @inbounds out[i, j, k] = a[i, j, k] - s / size(a, 3)
    end
end

function bench(Nx, Ny, Nz; n = 300)
    a = rand(Nx, Ny, Nz)
    out = similar(a)
    kernel = column_sum!(CPU())
    kernel(out, a; ndrange = (Nx, Ny))
    synchronize(CPU())
    return 1e6 * minimum(1:n) do _
        @elapsed begin
            kernel(out, a; ndrange = (Nx, Ny))
            synchronize(CPU())
        end
    end
end

for (Nx, Ny, Nz) in ((132, 132, 32), (100, 100, 50), (256, 256, 16))
    @printf("%2d threads  %3d×%3d×%2d  %7.1f µs\n", Threads.nthreads(), Nx, Ny, Nz, bench(Nx, Ny, Nz))
end

Tests

The "POCL workgroup size" testset now checks the new rule instead of PoCL's choices, with a regression case for the 132×132 launch. The "POCL bounds checks" testset launched 4097 work-items to force partial work-groups, but on 4 threads the new rule fits them into 17 work-groups of 241. It now uses a prime ndrange, which no work-group size divides.

[github-actions]

Documentation preview: https://juliagpu.github.io/KernelAbstractions.jl/previews/PR956/

[github-actions]

Benchmark results

Compared with 8a985320e, the latest measured ancestor of the merge base (14 commits before it).
The environment differs from the baseline's: Hwloc_jll 2.15.0+0 → 2.15.0+1, LinearAlgebraExt 0.10.5 → 0.1.46, StaticArraysCore 1.4.4 → 1.4.5.

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.

	Geometric mean	Unusual: slower	Unusual: faster	Compared
Performance	−19.6%	1	18	32
Latency	−56.3%	3	4	13
Applications	−25.8%	0	2	5

Benchmark	Baseline	This PR	Change
Performance
pocl/parallel/stream/triad/cache	7.50 µs	1.13 µs	−84.9%	🟢 faster
pocl/parallel/indexing/cartesian3d	83.7 µs	40.3 µs	−51.8%	🟢 faster
pocl/parallel/indexing/ntuple3d	82.6 µs	40.6 µs	−50.8%	🟢 faster
pocl/parallel/private/across-barrier	84.1 µs	42.2 µs	−49.8%	🟢 faster
pocl/parallel/indexing/linear3d	72.4 µs	40.3 µs	−44.4%	🟢 faster
pocl/parallel/localmem/transpose-tiled	110 µs	63.2 µs	−42.3%	🟢 faster
pocl/parallel/scheduling/default	181 µs	117 µs	−35.3%	🟢 faster
pocl/parallel/indexing/ntuple4d	168 µs	127 µs	−24.7%	🟢 faster
pocl/parallel/scheduling/explicit	154 µs	117 µs	−24.0%	🟢 faster
pocl/parallel/groupscan/add-i32	300 µs	392 µs	+30.8%	🔴 slower
pocl/serial/narrow/masked	797 µs	632 µs	−20.7%	🟢 faster
pocl/serial/math/sin-f64	210 µs	169 µs	−19.5%	🟢 faster
pocl/serial/narrow/unmasked	781 µs	632 µs	−19.1%	🟢 faster
pocl/parallel/util/zeros-large	996 µs	911 µs	−8.6%	🟢 faster
pocl/serial/arithmetic/muladd	765 µs	705 µs	−7.9%	🟢 faster
pocl/parallel/groupreduce/add-f64	1.17 ms	1.09 ms	−6.7%	🟢 faster
pocl/serial/const/marked	714 µs	668 µs	−6.5%	🟢 faster
pocl/serial/const/unmarked	914 µs	859 µs	−6.0%	🟢 faster
pocl/parallel/localmem/transpose-naive	457 µs	434 µs	−4.9%	🟢 faster
Latency
pocl/parallel/launch/dynamic-default	4.48 µs	165 ns	−96.3%	🟢 faster
pocl/parallel/launch/partial-workgroup	4.25 µs	218 ns	−94.9%	🟢 faster
pocl/parallel/launch/many-args	12.5 µs	656 ns	−94.8%	🟢 faster
pocl/parallel/launch/groupreduce-f64	6.53 µs	809 ns	−87.6%	🟢 faster
pocl/serial/first-kernel/warm-cache	275 ms	324 ms	+18.2%	🔴 slower
host/import/ka	501 ms	588 ms	+17.5%	🔴 slower
pocl/serial/first-kernel/cold-cache	353 ms	399 ms	+13.2%	🔴 slower
Applications
pocl/parallel/dgsem/weak-node-localmem	146 µs	83.0 µs	−43.3%	🟢 faster
pocl/serial/stencil/timestep	11.5 ms	9.85 ms	−14.0%	🟢 faster

All benchmarks

Benchmark	Baseline	This PR	Change
Performance
julia/serial/arithmetic/muladd-reference	655 µs	655 µs	+0.0%
julia/serial/math/exp-f64-reference	284 µs	291 µs	+2.5%
julia/serial/narrow/masked-reference	736 µs	732 µs	−0.6%
pocl/parallel/atomics/histogram	5.16 ms	5.21 ms	+0.9%
pocl/parallel/groupreduce/add-f64	1.17 ms	1.09 ms	−6.7%	🟢 faster
pocl/parallel/groupreduce/fallback	4.42 ms	4.38 ms	−0.8%
pocl/parallel/groupscan/add-i32	300 µs	392 µs	+30.8%	🔴 slower
pocl/parallel/indexing/cartesian3d	83.7 µs	40.3 µs	−51.8%	🟢 faster
pocl/parallel/indexing/linear3d	72.4 µs	40.3 µs	−44.4%	🟢 faster
pocl/parallel/indexing/ntuple3d	82.6 µs	40.6 µs	−50.8%	🟢 faster
pocl/parallel/indexing/ntuple4d	168 µs	127 µs	−24.7%	🟢 faster
pocl/parallel/localmem/transpose-naive	457 µs	434 µs	−4.9%	🟢 faster
pocl/parallel/localmem/transpose-tiled	110 µs	63.2 µs	−42.3%	🟢 faster
pocl/parallel/private/across-barrier	84.1 µs	42.2 µs	−49.8%	🟢 faster
pocl/parallel/scheduling/default	181 µs	117 µs	−35.3%	🟢 faster
pocl/parallel/scheduling/explicit	154 µs	117 µs	−24.0%	🟢 faster
pocl/parallel/stream/copy/cache	677 ns	720 ns	+6.4%
pocl/parallel/stream/copy/dram	6.52 ms	6.33 ms	−2.9%
pocl/parallel/stream/triad/cache	7.50 µs	1.13 µs	−84.9%	🟢 faster
pocl/parallel/stream/triad/dram	9.47 ms	9.27 ms	−2.2%
pocl/parallel/util/copyto-small	26.0 ns	25.0 ns	−3.8%
pocl/parallel/util/zeros-large	996 µs	911 µs	−8.6%	🟢 faster
pocl/parallel/util/zeros-small	643 ns	647 ns	+0.6%
pocl/serial/arithmetic/contract	1.90 ms	1.85 ms	−2.7%
pocl/serial/arithmetic/muladd	765 µs	705 µs	−7.9%	🟢 faster
pocl/serial/const/marked	714 µs	668 µs	−6.5%	🟢 faster
pocl/serial/const/unmarked	914 µs	859 µs	−6.0%	🟢 faster
pocl/serial/math/exp-f64	139 µs	141 µs	+2.0%
pocl/serial/math/sin-f64	210 µs	169 µs	−19.5%	🟢 faster
pocl/serial/narrow/masked	797 µs	632 µs	−20.7%	🟢 faster
pocl/serial/narrow/unmasked	781 µs	632 µs	−19.1%	🟢 faster
pocl/serial/stream/triad/dram	8.67 ms	8.50 ms	−1.9%
Latency
host/import/ka	501 ms	588 ms	+17.5%	🔴 slower
host/precompile/ka	39.7 s	41.4 s	+4.4%
host/startup	131 ms	131 ms	−0.2%
pocl/parallel/launch/dynamic-default	4.48 µs	165 ns	−96.3%	🟢 faster
pocl/parallel/launch/groupreduce-f64	6.53 µs	809 ns	−87.6%	🟢 faster
pocl/parallel/launch/many-args	12.5 µs	656 ns	−94.8%	🟢 faster
pocl/parallel/launch/partial-workgroup	4.25 µs	218 ns	−94.9%	🟢 faster
pocl/serial/compile/groupreduce	1.49 s	1.48 s	−0.2%
pocl/serial/compile/localmem-barrier	710 ms	727 ms	+2.4%
pocl/serial/compile/new-kernel	121 ms	128 ms	+5.8%
pocl/serial/compile/new-workgroupsize	8.79 ms	8.68 ms	−1.3%
pocl/serial/first-kernel/cold-cache	353 ms	399 ms	+13.2%	🔴 slower
pocl/serial/first-kernel/warm-cache	275 ms	324 ms	+18.2%	🔴 slower
Applications
pocl/parallel/dgsem/fluxdiff-node	289 µs	253 µs	−12.7%
pocl/parallel/dgsem/weak-element	161 µs	93.1 µs	−42.1%
pocl/parallel/dgsem/weak-node-localmem	146 µs	83.0 µs	−43.3%	🟢 faster
pocl/parallel/stencil/timestep	8.11 ms	7.38 ms	−9.0%
pocl/serial/stencil/timestep	11.5 ms	9.85 ms	−14.0%	🟢 faster

History: https://juliagpu.org/perf/?pkg=kernelabstractions

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.81%. Comparing base (cc7cf19) to head (d9b042c).

Additional details and impacted files

@@                     Coverage Diff                     @@
##           tb/pocl-launch-cost-906     #956      +/-   ##
===========================================================
- Coverage                    84.81%   84.81%   -0.01%     
===========================================================
  Files                           31       31              
  Lines                         2924     2923       -1     
===========================================================
- Hits                          2480     2479       -1     
  Misses                         444      444              

☔ 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.