Fix zeros erroring on an empty shape (#515)

JuliaGPU · Aug 27, 2024 · 5f99533 · 5f99533
1 parent 05a3c05
commit 5f99533
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/cpu.jl b/src/cpu.jl
@@ -49,7 +49,11 @@ 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
+    n = prod(ndrange)
+    if iszero(n)
+        # If the ndrange is zero return a workgroupsize of (1, 1,...)
+        return map(one, ndrange)
+    elseif n <= CPU_GRAINSIZE
         return ndrange
     else
         available = Ref(CPU_GRAINSIZE)

diff --git a/test/test.jl b/test/test.jl
@@ -296,4 +296,14 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk
         @test KernelAbstractions.default_cpu_workgroupsize((5, 7, 13, 17)) == (5, 7, 13, 2)
     end
 
+    @testset "empty arrays" begin
+        backend = Backend()
+        @test size(allocate(backend, Float32, 0)) == (0,)
+        @test size(allocate(backend, Float32, 3, 0)) == (3, 0)
+        @test size(allocate(backend, Float32, 0, 9)) == (0, 9)
+        @test size(KernelAbstractions.zeros(backend, Float32, 0)) == (0,)
+        @test size(KernelAbstractions.zeros(backend, Float32, 3, 0)) == (3, 0)
+        @test size(KernelAbstractions.zeros(backend, Float32, 0, 9)) == (0, 9)
+    end
+
 end