Implement groupreduce API

docs/src/api.md

@@ -13,13 +13,14 @@
 @uniform
 @groupsize
 @ndrange
-synchronize
-allocate
+@groupreduce
 ```
 
 ## Host language
 
 ```@docs
+synchronize
+allocate
 KernelAbstractions.zeros
 ```
 

src/KernelAbstractions.jl

@@ -798,6 +798,8 @@ function __fake_compiler_job end
 # - LoopInfo
 ###
 
+include("reduce.jl")
+
 include("extras/extras.jl")
 
 include("reflection.jl")

src/reduce.jl

@@ -0,0 +1,130 @@
+export @groupreduce, Reduction
+
+module Reduction
+    const thread = Val(:thread)
+    const warp = Val(:warp)
+end
+
+"""
+    @groupreduce op val neutral algo [groupsize]
+
+Perform group reduction of `val` using `op`.
+
+# 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.
+    If a kernel does not specifies `groupsize` statically, then it is required to
+    provide `groupsize`.
+    Also can be used to perform reduction accross first `groupsize` threads
+    (if `groupsize < @groupsize()`).
+
+# Returns
+
+Result of the reduction.
+"""
+macro groupreduce(op, val, neutral, algo)
+    return quote
+        __groupreduce(
+            $(esc(:__ctx__)),
+            $(esc(op)),
+            $(esc(val)),
+            $(esc(neutral)),
+            Val(prod($groupsize($(esc(:__ctx__))))),
+            $(esc(algo)),
+        )
+    end
+end
+
+macro groupreduce(op, val, neutral, algo, groupsize)
+    return quote
+        __groupreduce(
+            $(esc(:__ctx__)),
+            $(esc(op)),
+            $(esc(val)),
+            $(esc(neutral)),
+            Val($(esc(groupsize))),
+            $(esc(algo)),
+        )
+    end
+end
+
+function __groupreduce(__ctx__, op, val::T, neutral::T, ::Val{groupsize}, ::Val{:thread}) where {T, groupsize}
+    storage = @localmem T groupsize
+
+    local_idx = @index(Local)
+    @inbounds local_idx ≤ groupsize && (storage[local_idx] = val)
+    @synchronize()
+
+    s::UInt64 = groupsize ÷ 0x02
+    while s > 0x00
+        if (local_idx - 0x01) < s
+            other_idx = local_idx + s
+            if other_idx ≤ groupsize
+                @inbounds storage[local_idx] = op(storage[local_idx], storage[other_idx])
+            end
+        end
+        @synchronize()
+        s >>= 0x01
+    end
+
+    if local_idx == 0x01
+        @inbounds val = storage[local_idx]
+    end
+    return val
+end
+
+# Warp groupreduce.
+
+macro shfl_down(val, offset)
+    return quote
+        $__shfl_down($(esc(val)), $(esc(offset)))
+    end
+end
+
+# Backends should implement these two.
+function __shfl_down end
+supports_warp_reduction(::Backend) = false
+
+# Assume warp is 32 lanes.
+const __warpsize = UInt32(32)
+# Maximum number of warps (for a groupsize = 1024).
+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))
+        offset >>= 0x01
+    end
+    return val
+end
+
+function __groupreduce(__ctx__, op, val::T, neutral::T, ::Val{groupsize}, ::Val{:warp}) where {T, groupsize}
+    storage = @localmem T __warp_bins
+
+    local_idx = @index(Local)
+    lane = (local_idx - 0x01) % __warpsize + 0x01
+    warp_id = (local_idx - 0x01) ÷ __warpsize + 0x01
+
+    # Each warp performs a reduction and writes results into its own bin in `storage`.
+    val = __warp_reduce(val, op)
+    @inbounds lane == 0x01 && (storage[warp_id] = val)
+    @synchronize()
+
+    # Final reduction of the `storage` on the first warp.
+    within_storage = (local_idx - 0x01) < groupsize ÷ __warpsize
+    @inbounds val = within_storage ? storage[lane] : neutral
+    warp_id == 0x01 && (val = __warp_reduce(val, op))
+    return val
+end

test/groupreduce.jl

@@ -0,0 +1,48 @@
+@kernel function groupreduce_1!(y, x, op, neutral, algo)
+    i = @index(Global)
+    val = i > length(x) ? neutral : x[i]
+    res = @groupreduce(op, val, neutral, algo)
+    i == 1 && (y[1] = res)
+end
+
+@kernel function groupreduce_2!(y, x, op, neutral, algo, ::Val{groupsize}) where {groupsize}
+    i = @index(Global)
+    val = i > length(x) ? neutral : x[i]
+    res = @groupreduce(op, val, neutral, algo, groupsize)
+    i == 1 && (y[1] = res)
+end
+
+function groupreduce_testsuite(backend, AT)
+    @testset "@groupreduce" begin
+        @testset "thread reduction T=$T, n=$n" for T in (Float16, Float32, Int32, Int64), n in (256, 512, 1024)
+            x = AT(ones(T, n))
+            y = AT(zeros(T, 1))
+
+            groupreduce_1!(backend(), n)(y, x, +, zero(T), Reduction.thread; ndrange = n)
+            @test Array(y)[1] == n
+
+            groupreduce_2!(backend())(y, x, +, zero(T), Reduction.thread, Val(128); ndrange = n)
+            @test Array(y)[1] == 128
+
+            groupreduce_2!(backend())(y, x, +, zero(T), Reduction.thread, 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 (256, 512, 1024)
+
+                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

test/testsuite.jl

@@ -38,6 +38,7 @@ include("reflection.jl")
 include("examples.jl")
 include("convert.jl")
 include("specialfunctions.jl")
+include("groupreduce.jl")
 
 function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}())
     @conditional_testset "Unittests" skip_tests begin
@@ -92,6 +93,13 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{
         examples_testsuite(backend_str)
     end
 
+    # TODO @index(Local) only works as a top-level expression on CPU.
+    if backend != CPU
+        @conditional_testset "@groupreduce" skip_tests begin
+            groupreduce_testsuite(backend, AT)
+        end
+    end
+
     return
 end