From 5d0e02c8f414c9758a4ec4df096c635ee4618ecc Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 09:28:22 +0100 Subject: [PATCH 1/2] POCL: shrink tuned work-groups to divide the ndrange MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A launch with a partial work-group runs the bounds-checked kernel, whose work-item loop vectorizes poorly. Shrink each dimension of the tuned size to the largest divisor of the ndrange that keeps at least half of it. On an M1, a 126³ stencil takes 257 µs instead of 588. --- src/pocl/backend.jl | 29 +++++++++++++++++++++++++++++ test/runtests.jl | 29 +++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 7ee0823c5..c3b9e68b1 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -257,6 +257,35 @@ function KA.expected_workgroupsize(::POCLBackend, nitems) ) end +# a launch with a partial work-group runs the bounds-checked kernel, whose work-item loop +# vectorizes poorly (AArch64, without masked loads and stores: ~3x slower for a stencil). +# so shrink each dimension of the generic choice to the largest divisor of the ndrange that +# keeps at least half its size, if there is one. +# (only shrink, so padding stays within the bound `KA.select_launch` assumes.) +function KA.launch_workgroupsize(::POCLBackend, ::KA.NDLaunch, threads, ndrange) + extent = KA.extents(ndrange) + limits = KA.nd_limits(device_limits().max_work_group_dims, extent) + items = KI.threads_to_workgroupsize(threads, extent, limits) + return map((n, x) -> n % x == 0 ? x : largest_divisor(n, x), extent, items) +end + +# the largest divisor of `n` in `cld(x, 2):x`, or `x` if there is none. searches the sizes +# or the work-group counts, whichever are fewer. requires `n % x != 0`: the size search +# skips `x`. +function largest_divisor(n::Int, x::Int) + lo = cld(x, 2) + if x - lo <= fld(n, lo) - cld(n, x) + for d in (x - 1):-1:lo + n % d == 0 && return d + end + else + for g in cld(n, x):fld(n, lo) + n % g == 0 && return n ÷ g + end + end + return x +end + # PoCL's own choice when a launch doesn't give a local size (`pocl_default_local_size_optimizer` # in `lib/CL/devices/pocl_local_size.c`), in one dimension: start from the largest work-group, # and split it while there are at most half as many work-groups as threads, but not below diff --git a/test/runtests.jl b/test/runtests.jl index 6f18d7481..95f8776de 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -917,6 +917,35 @@ end @test KernelAbstractions.select_launch(kernel, wgs, iterspace) === KernelAbstractions.NDLaunch{Int}() end +@testset "POCL N-d work-group sizes" begin + # tuned work-groups shrink to divisors of the ndrange where possible, avoiding partial groups + wgs(ndrange) = KernelAbstractions.launch_workgroupsize( + CPU(), KernelAbstractions.NDLaunch{Int}(), 4096, ndrange + ) + # (searching the work-group counts) + @test wgs((126, 126, 126)) == (126, 21, 1) + @test wgs((100, 100)) == (100, 25) + @test wgs((5000, 10)) == (2500, 1) + # (searching the sizes) + @test wgs((1024, 1002)) == (1024, 3) + @test wgs((64, 64, 64)) == (64, 64, 1) + # each dimension keeps at least half its size + @test wgs((127, 127, 127)) == (127, 32, 1) + @test wgs((1024, 1001)) == (1024, 4) + @test wgs((0, 7)) == (1, 7) + + # an explicit workgroup size is kept + @kernel function groupsize!(A) + if @index(Global, Linear) == 1 + A[1], A[2] = @groupsize() + end + end + A = zeros(Int, 2) + groupsize!(CPU())(A; ndrange = (126, 126), workgroupsize = (126, 32)) + KernelAbstractions.synchronize(CPU()) + @test A == [126, 32] +end + @kernel function fill_index!(A) I = @index(Global, Linear) @inbounds A[I] = I From 9093debb009ccf1f6681610428e58315dba4880a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 28 Oct 2026 09:28:31 +0100 Subject: [PATCH 2/2] POCL: shrink tuned work-groups to a quarter when no divisor fits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Such launches run the bounds-checked kernel anyway, and run it faster in smaller work-groups: axpy over 1000003 elements on 4 threads takes 31 µs instead of 71 on a Ryzen 9 9950X, and 128 instead of 146 on an EPYC 7502. --- src/pocl/backend.jl | 18 +++++++++++++----- test/runtests.jl | 7 ++++--- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index c3b9e68b1..21e7914cd 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -260,18 +260,26 @@ end # a launch with a partial work-group runs the bounds-checked kernel, whose work-item loop # vectorizes poorly (AArch64, without masked loads and stores: ~3x slower for a stencil). # so shrink each dimension of the generic choice to the largest divisor of the ndrange that -# keeps at least half its size, if there is one. +# keeps at least half its size. without one, the launch stays checked, but runs faster in +# work-groups a quarter of the size (axpy, 1000003 work-items on 4 threads of a Ryzen 9 +# 9950X: 71 vs 31 µs). # (only shrink, so padding stays within the bound `KA.select_launch` assumes.) function KA.launch_workgroupsize(::POCLBackend, ::KA.NDLaunch, threads, ndrange) extent = KA.extents(ndrange) limits = KA.nd_limits(device_limits().max_work_group_dims, extent) items = KI.threads_to_workgroupsize(threads, extent, limits) - return map((n, x) -> n % x == 0 ? x : largest_divisor(n, x), extent, items) + return map(fit_workgroup, extent, items) end -# the largest divisor of `n` in `cld(x, 2):x`, or `x` if there is none. searches the sizes +function fit_workgroup(n::Int, x::Int) + n % x == 0 && return x + d = largest_divisor(n, x) + return d == 0 ? cld(x, 4) : d +end + +# the largest divisor of `n` in `cld(x, 2):x`, or 0 if there is none. searches the sizes # or the work-group counts, whichever are fewer. requires `n % x != 0`: the size search -# skips `x`. +# skips `x` (`largest_divisor(8, 4) == 2`). function largest_divisor(n::Int, x::Int) lo = cld(x, 2) if x - lo <= fld(n, lo) - cld(n, x) @@ -283,7 +291,7 @@ function largest_divisor(n::Int, x::Int) n % g == 0 && return n ÷ g end end - return x + return 0 end # PoCL's own choice when a launch doesn't give a local size (`pocl_default_local_size_optimizer` diff --git a/test/runtests.jl b/test/runtests.jl index 95f8776de..1c94d8cfb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -929,9 +929,10 @@ end # (searching the sizes) @test wgs((1024, 1002)) == (1024, 3) @test wgs((64, 64, 64)) == (64, 64, 1) - # each dimension keeps at least half its size - @test wgs((127, 127, 127)) == (127, 32, 1) - @test wgs((1024, 1001)) == (1024, 4) + # each dimension keeps at least half its size, or shrinks to a quarter without a divisor + @test wgs((127, 127, 127)) == (127, 8, 1) + @test wgs((1024, 1001)) == (1024, 1) + @test wgs((100003,)) == (1024,) @test wgs((0, 7)) == (1, 7) # an explicit workgroup size is kept