Auto-select reduction algorithm & remove at-shfl_down macro

Author: pxl-th

src/reduce.jl

@@ -1,26 +1,18 @@
-export @groupreduce, Reduction
+export @groupreduce
 
 module Reduction
     const thread = Val(:thread)
     const warp = Val(:warp)
 end
 
 """
-    @groupreduce op val neutral algo [groupsize]
+    @groupreduce op val neutral [groupsize]
 
 Perform group reduction of `val` using `op`.
+If backend supports warp reduction, it will use it instead of thread reduction.
 
 # Arguments
 
-- `algo` specifies which reduction algorithm to use:
-    - `Reduction.thread`:
-        Perform thread group reduction (requires `groupsize * sizeof(T)` bytes of shared memory).
-        Available accross all backends.
-    - `Reduction.warp`:
-        Perform warp group reduction (requires `32 * sizeof(T)` bytes of shared memory).
-        Potentially faster, since requires fewer writes to shared memory.
-        To query if backend supports warp reduction, use `supports_warp_reduction(backend)`.
-
 - `neutral` should be a neutral w.r.t. `op`, such that `op(neutral, x) == x`.
 
 - `groupsize` specifies size of the workgroup.
@@ -33,29 +25,51 @@ Perform group reduction of `val` using `op`.
 
 Result of the reduction.
 """
-macro groupreduce(op, val, neutral, algo)
+macro groupreduce(op, val, neutral)
     return quote
-        __groupreduce(
-            $(esc(:__ctx__)),
-            $(esc(op)),
-            $(esc(val)),
-            $(esc(neutral)),
-            Val(prod($groupsize($(esc(:__ctx__))))),
-            $(esc(algo)),
-        )
+        if __supports_warp_reduction()
+            __groupreduce(
+                $(esc(:__ctx__)),
+                $(esc(op)),
+                $(esc(val)),
+                $(esc(neutral)),
+                Val(prod($groupsize($(esc(:__ctx__))))),
+                $(esc(Reduction.warp)),
+            )
+        else
+            __groupreduce(
+                $(esc(:__ctx__)),
+                $(esc(op)),
+                $(esc(val)),
+                $(esc(neutral)),
+                Val(prod($groupsize($(esc(:__ctx__))))),
+                $(esc(Reduction.thread)),
+            )
+        end
     end
 end
 
-macro groupreduce(op, val, neutral, algo, groupsize)
+macro groupreduce(op, val, neutral, groupsize)
     return quote
-        __groupreduce(
-            $(esc(:__ctx__)),
-            $(esc(op)),
-            $(esc(val)),
-            $(esc(neutral)),
-            Val($(esc(groupsize))),
-            $(esc(algo)),
-        )
+        if __supports_warp_reduction()
+            __groupreduce(
+                $(esc(:__ctx__)),
+                $(esc(op)),
+                $(esc(val)),
+                $(esc(neutral)),
+                Val($(esc(groupsize))),
+                $(esc(Reduction.warp)),
+            )
+        else
+            __groupreduce(
+                $(esc(:__ctx__)),
+                $(esc(op)),
+                $(esc(val)),
+                $(esc(neutral)),
+                Val($(esc(groupsize))),
+                $(esc(Reduction.thread)),
+            )
+        end
     end
 end
 
@@ -86,15 +100,9 @@ end
 
 # Warp groupreduce.
 
-macro shfl_down(val, offset)
-    return quote
-        $__shfl_down($(esc(val)), $(esc(offset)))
-    end
-end
-
-# Backends should implement these two.
+# NOTE: Backends should implement these two device functions (with `@device_override`).
 function __shfl_down end
-supports_warp_reduction(::Backend) = false
+function __supports_warp_reduction() end
 
 # Assume warp is 32 lanes.
 const __warpsize = UInt32(32)
@@ -104,7 +112,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,51 +1,35 @@
-@kernel function groupreduce_1!(y, x, op, neutral, algo)
+@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, algo)
+    res = @groupreduce(op, val, neutral)
     i == 1 && (y[1] = res)
 end
 
-@kernel function groupreduce_2!(y, x, op, neutral, algo, ::Val{groupsize}) where {groupsize}
+@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, algo, groupsize)
+    res = @groupreduce(op, val, neutral, groupsize)
     i == 1 && (y[1] = res)
 end
 
 function groupreduce_testsuite(backend, AT)
-    # TODO should be better way of querying max groupsize
+    # TODO should be a better way of querying max groupsize
     groupsizes = "$backend" == "oneAPIBackend" ?
         (256,) :
         (256, 512, 1024)
     @testset "@groupreduce" begin
-        @testset "thread reduction T=$T, n=$n" for T in (Float16, Float32, Int32, Int64), n in groupsizes
+        @testset "T=$T, n=$n" for T in (Float16, Float32, Float64, Int16, Int32, Int64), n in groupsizes
             x = AT(ones(T, n))
             y = AT(zeros(T, 1))
 
-            groupreduce_1!(backend(), n)(y, x, +, zero(T), Reduction.thread; ndrange = n)
+            groupreduce_1!(backend(), n)(y, x, +, zero(T); ndrange = n)
             @test Array(y)[1] == n
 
-            groupreduce_2!(backend())(y, x, +, zero(T), Reduction.thread, Val(128); ndrange = n)
+            groupreduce_2!(backend())(y, x, +, zero(T), Val(128); ndrange = n)
             @test Array(y)[1] == 128
 
-            groupreduce_2!(backend())(y, x, +, zero(T), Reduction.thread, Val(64); ndrange = n)
+            groupreduce_2!(backend())(y, x, +, zero(T), Val(64); ndrange = n)
             @test Array(y)[1] == 64
         end
-
-        warp_reduction = KernelAbstractions.supports_warp_reduction(backend())
-        if warp_reduction
-            @testset "warp reduction T=$T, n=$n" for T in (Float16, Float32, Int32, Int64), n in groupsizes
-                x = AT(ones(T, n))
-                y = AT(zeros(T, 1))
-                groupreduce_1!(backend(), n)(y, x, +, zero(T), Reduction.warp; ndrange = n)
-                @test Array(y)[1] == n
-
-                groupreduce_2!(backend())(y, x, +, zero(T), Reduction.warp, Val(128); ndrange = n)
-                @test Array(y)[1] == 128
-
-                groupreduce_2!(backend())(y, x, +, zero(T), Reduction.warp, Val(64); ndrange = n)
-                @test Array(y)[1] == 64
-            end
-        end
     end
 end