Repository: JuliaGPU/KernelAbstractions.jl (julia)

Issue:
Default CPU workgroupsize can be inadequat for higher-dimensionsal kernels
As observed in https://github.com/WaterLily-jl/WaterLily.jl/pull/133#issuecomment-2263500896

When we have a grid like `(64, 64, 64)` using a WG size of (1024, 1, 1) is inefficient. Instead WL uses `(64,1,1)` and maybe KA should "spread" the `1024` to `(64, 64,1)`.

Pull request description:
Fixes #499

cc: @b-fg @weymouth

Patch:
diff --git a/src/cpu.jl b/src/cpu.jl
index 151b31f2..539a1e33 100644
--- a/src/cpu.jl
+++ b/src/cpu.jl
@@ -46,6 +46,29 @@ function (obj::Kernel{CPU})(args...; ndrange=nothing, workgroupsize=nothing, )
     __run(obj, ndrange, iterspace, args, dynamic, obj.backend.static)
 end
 
+const CPU_GRAINSIZE = 1024 # Vectorization, 4x unrolling, minimal grain size
+function default_cpu_workgroupsize(ndrange)
+    # if the total kernel is small, don't launch multiple tasks
+    if prod(ndrange) <= CPU_GRAINSIZE
+        return ndrange
+    else
+        available = Ref(CPU_GRAINSIZE)
+        return ntuple(length(ndrange)) do i
+            dim = ndrange[i]
+            remaining = available[]
+            if remaining == 0
+                return 1
+            elseif remaining <= dim
+                available[] = 0
+                return remaining
+            else
+                available[] = remaining ÷ dim
+                return dim
+            end
+        end
+    end
+end
+
 @inline function launch_config(kernel::Kernel{CPU}, ndrange, workgroupsize)
     if ndrange isa Integer
         ndrange = (ndrange,)
@@ -55,7 +78,7 @@ end
     end
 
     if KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing
-        workgroupsize = (1024,) # Vectorization, 4x unrolling, minimal grain size
+        workgroupsize = default_cpu_workgroupsize(ndrange)
     end
     iterspace, dynamic = partition(kernel, ndrange, workgroupsize)
     # partition checked that the ndrange's agreed


Test patch:
diff --git a/test/test.jl b/test/test.jl
index 2eda66e0..9bf6f849 100644
--- a/test/test.jl
+++ b/test/test.jl
@@ -282,4 +282,14 @@ end
     @test KernelAbstractions.functional(Backend()) isa Union{Missing, Bool}
 end
 
+@testset "CPU default workgroupsize" begin
+    @test KernelAbstractions.default_cpu_workgroupsize((64,)) == (64,)
+    @test KernelAbstractions.default_cpu_workgroupsize((1024,)) == (1024,)
+    @test KernelAbstractions.default_cpu_workgroupsize((2056,)) == (1024,)
+    @test KernelAbstractions.default_cpu_workgroupsize((64, 64,)) == (64,16)
+    @test KernelAbstractions.default_cpu_workgroupsize((64, 64, 64,4)) == (64,16,1,1)
+    @test KernelAbstractions.default_cpu_workgroupsize((64,15)) == (64,15)
+    @test KernelAbstractions.default_cpu_workgroupsize((5,7,13,17)) == (5,7,13,2)
+end
+
 end