diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 7ee0823c5..21e7914cd 100644 --- a/src/pocl/backend.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 6f18d7481..1c94d8cfb 100644 --- a/test/runtests.jl +++ b/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