groupreduction and subgroupreduction

src/KernelAbstractions.jl

@@ -155,10 +155,10 @@ function unsafe_free! end
 # - @groupsize
 # - @ndrange
 ###
-
 function groupsize end
 function ndrange end
 
+
 """
     @groupsize()
 
@@ -676,6 +676,7 @@ function __synchronize()
     error("@synchronize used outside kernel or not captured")
 end
 
+
 @generated function __print(items...)
     str = ""
     args = []
@@ -719,6 +720,7 @@ end
     @inbounds A[I] = B[I]
 end
 
+
 # CPU backend
 
 include("cpu.jl")
@@ -745,4 +747,7 @@ end
     end
 end
 
+# groupreduce
+include("reduce.jl")
+
 end #module

src/reduce.jl

@@ -0,0 +1,53 @@
+export @groupreduce
+
+"""
+    @groupreduce(op, val, neutral, use_subgroups)
+
+Reduce values across a block
+- `op`: the operator of the reduction
+- `val`: value that each thread contibutes to the values that need to be reduced
+- `neutral`: value of the operator, so that `op(netural, neutral) = neutral``
+- `groupsize` (optional): specify the groupszie. If not specified @groupsize is used but this is generally slower.
+"""
+macro groupreduce(op, val, neutral) 
+  quote
+      $__groupreduce($(esc(:__ctx__)),$(esc(op)), $(esc(val)), $(esc(neutral)), Val(prod(groupsize($(esc(:__ctx__))))))
+  end
+end
+
+macro groupreduce(op, val, neutral, groupsize) 
+    quote
+        $__groupreduce($(esc(:__ctx__)),$(esc(op)), $(esc(val)), $(esc(neutral)), $(esc(groupsize)))
+    end
+end
+
+@inline function __groupreduce(__ctx__, op, val::T, neutral, ::Val{GROUPSIZE}) where {T, GROUPSIZE}
+    idx_in_group = @index(Local)
+
+    localmem = @localmem(T, GROUPSIZE)
+
+    @inbounds localmem[idx_in_group] = val
+
+    # perform the reduction
+    d = 1
+    while d < GROUPSIZE
+        @synchronize()
+        index = 2 * d * (idx_in_group-1) + 1
+        @inbounds if index <= GROUPSIZE
+            other_val = if index + d <= GROUPSIZE
+                localmem[index+d]
+            else
+                neutral
+            end
+            localmem[index] = op(localmem[index], other_val)
+        end
+        d *= 2
+    end
+
+    # load the final value on the first thread
+    if idx_in_group == 1
+        val = @inbounds localmem[idx_in_group]
+    end
+
+    return val 
+end

test/reduce.jl

@@ -0,0 +1,42 @@
+using KernelAbstractions, Test
+
+@kernel function reduce(a, b, op, neutral)
+    I = @index(Global)
+    gI = @index(Group, Linear)
+    val = a[I]
+
+    val = @groupreduce(op, val, neutral)
+
+    b[gI] = val 
+end
+
+function reduce_testset(backend, ArrayT)
+    @testset "groupreduce one group" begin
+        @testset for op in (+, *, max, min)
+            @testset for type in (Int32, Float32, Float64)
+                @test test_groupreduce(backend, ArrayT, op, type, op(neutral), 8)
+                @test test_groupreduce(backend, ArrayT, op, type, op(neutral), 16)
+                @test test_groupreduce(backend, ArrayT, op, type, op(neutral), 32)
+                @test test_groupreduce(backend, ArrayT, op, type, op(neutral), 64)
+            end
+        end
+    end
+end
+
+function test_groupreduce(backend, ArrayT, op, type, neutral, N)
+    a = rand(type, N)
+    b = ArrayT(a)
+
+    gsz = 64
+    ngroups = ceil(N/gsz)
+    c = similar(b, ngroups)
+    kernel = reduce(backend, (gsz,))
+    kernel(a, c, op, neutral)
+
+    expected = mapreduce(x->x^2, +, a)
+    actual = c[1]
+    return expected = actual
+end
+
+
+

test/testsuite.jl

@@ -32,6 +32,7 @@ include("reflection.jl")
 include("examples.jl")
 include("convert.jl")
 include("specialfunctions.jl")
+include("reduce.jl")
 
 function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}())
     @conditional_testset "Unittests" skip_tests begin
@@ -78,6 +79,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{
         convert_testsuite(backend, AT)
     end
 
+    @conditional_testset "Reduce" skip_tests begin
+        reduce_testsuite(backend, AT)
+    end
+
     @conditional_testset "Examples" skip_tests begin
         examples_testsuite(backend_str)
     end