POCL: fit tuned work-groups to the ndrange

src/pocl/backend.jl

@@ -257,6 +257,43 @@ 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. 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(fit_workgroup, extent, items)
+end
+
+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` (`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)
+        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 0
+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

test/runtests.jl

@@ -917,6 +917,36 @@ 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, 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
+    @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