Add warp_groupreduce tests

Author: pxl-th

src/reduce.jl

@@ -62,10 +62,10 @@ end
 # Warp groupreduce.
 
 # NOTE: Backends should implement these two device functions (with `@device_override`).
-function __shfl_down end
-function __supports_warp_reduction()
-    return false
-end
+function shfl_down end
+supports_warp_reduction() = false
+# Host-variant.
+supports_warp_reduction(::Backend) = false
 
 # Assume warp is 32 lanes.
 const __warpsize = UInt32(32)
@@ -75,7 +75,7 @@ const __warp_bins = UInt32(32)
 @inline function __warp_reduce(val, op)
     offset::UInt32 = __warpsize ÷ 0x02
     while offset > 0x00
-        val = op(val, __shfl_down(val, offset))
+        val = op(val, shfl_down(val, offset))
         offset >>= 0x01
     end
     return val

test/groupreduce.jl

@@ -1,14 +1,28 @@
 @kernel cpu=false function groupreduce_1!(y, x, op, neutral)
     i = @index(Global)
     val = i > length(x) ? neutral : x[i]
-    res = @groupreduce(op, val, neutral)
+    res = @groupreduce(op, val)
     i == 1 && (y[1] = res)
 end
 
 @kernel cpu=false function groupreduce_2!(y, x, op, neutral, ::Val{groupsize}) where {groupsize}
     i = @index(Global)
     val = i > length(x) ? neutral : x[i]
-    res = @groupreduce(op, val, neutral, groupsize)
+    res = @groupreduce(op, val, groupsize)
+    i == 1 && (y[1] = res)
+end
+
+@kernel cpu=false function warp_groupreduce_1!(y, x, op, neutral)
+    i = @index(Global)
+    val = i > length(x) ? neutral : x[i]
+    res = @warp_groupreduce(op, val, neutral)
+    i == 1 && (y[1] = res)
+end
+
+@kernel cpu=false function warp_groupreduce_2!(y, x, op, neutral, ::Val{groupsize}) where {groupsize}
+    i = @index(Global)
+    val = i > length(x) ? neutral : x[i]
+    res = @warp_groupreduce(op, val, neutral, groupsize)
     i == 1 && (y[1] = res)
 end
 
@@ -17,19 +31,40 @@ function groupreduce_testsuite(backend, AT)
     groupsizes = "$backend" == "oneAPIBackend" ?
         (256,) :
         (256, 512, 1024)
+
     @testset "@groupreduce" begin
         @testset "T=$T, n=$n" for T in (Float16, Float32, Int16, Int32, Int64), n in groupsizes
             x = AT(ones(T, n))
             y = AT(zeros(T, 1))
+            neutral = zero(T)
+            op = +
 
-            groupreduce_1!(backend(), n)(y, x, +, zero(T); ndrange = n)
+            groupreduce_1!(backend(), n)(y, x, op, neutral; ndrange = n)
             @test Array(y)[1] == n
 
-            groupreduce_2!(backend())(y, x, +, zero(T), Val(128); ndrange = n)
-            @test Array(y)[1] == 128
+            for groupsize in (64, 128)
+                groupreduce_2!(backend())(y, x, op, neutral, Val(groupsize); ndrange = n)
+                @test Array(y)[1] == groupsize
+            end
+        end
+    end
+
+    if KernelAbstractions.supports_warp_reduction(backend())
+        @testset "@warp_groupreduce" begin
+            @testset "T=$T, n=$n" for T in (Float16, Float32, Int16, Int32, Int64), n in groupsizes
+                x = AT(ones(T, n))
+                y = AT(zeros(T, 1))
+                neutral = zero(T)
+                op = +
+
+                warp_groupreduce_1!(backend(), n)(y, x, op, neutral; ndrange = n)
+                @test Array(y)[1] == n
 
-            groupreduce_2!(backend())(y, x, +, zero(T), Val(64); ndrange = n)
-            @test Array(y)[1] == 64
+                for groupsize in (64, 128)
+                    warp_groupreduce_2!(backend())(y, x, op, neutral, Val(groupsize); ndrange = n)
+                    @test Array(y)[1] == groupsize
+                end
+            end
         end
     end
 end