Simplify algo selection

Author: pxl-th

docs/src/api.md

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

src/reduce.jl

@@ -1,18 +1,28 @@
+export @groupreduce, Reduction
+
+module Reduction
+    const thread = Val(:thread)
+    const warp = Val(:warp)
+end
+
 """
-    @groupreduce algo op val neutral [groupsize]
+    @groupreduce op val neutral algo [groupsize]
 
 Perform group reduction of `val` using `op`.
 
 # Arguments
 
 - `algo` specifies which reduction algorithm to use:
-    - `:thread`:
+    - `Reduction.thread`:
         Perform thread group reduction (requires `groupsize * sizeof(T)` bytes of shared memory).
         Available accross all backends.
-    - `:warp`:
+    - `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`.
@@ -23,69 +33,53 @@ Perform group reduction of `val` using `op`.
 
 Result of the reduction.
 """
-macro groupreduce(algo, op, val, neutral)
-    f = if algo.value == :thread
-        __groupreduce
-    elseif algo.value == :warp
-        __warp_groupreduce
-    else
-        error(
-            "@groupreduce supports only :thread or :warp as a reduction algorithm, " *
-            "but $(algo.value) was specified.")
-    end
+macro groupreduce(op, val, neutral, algo)
     quote
-        $f(
+        __groupreduce(
             $(esc(:__ctx__)),
             $(esc(op)),
             $(esc(val)),
             $(esc(neutral)),
             Val(prod($groupsize($(esc(:__ctx__))))),
+            $(esc(algo)),
         )
     end
 end
 
-macro groupreduce(algo, op, val, neutral, groupsize)
-    f = if algo.value == :thread
-        __groupreduce
-    elseif algo.value == :warp
-        __warp_groupreduce
-    else
-        error(
-            "@groupreduce supports only :thread or :warp as a reduction algorithm, " *
-            "but $(algo.value) was specified.")
-    end
+macro groupreduce(op, val, neutral, algo, groupsize)
     quote
-        $f(
+        __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}) where {T, groupsize}
+function __groupreduce(__ctx__, op, val::T, neutral::T, ::Val{groupsize}, ::Val{:thread}) where {T, groupsize}
     storage = @localmem T groupsize
 
     local_idx = @index(Local)
-    local_idx ≤ groupsize && (storage[local_idx] = val)
+    @inbounds local_idx ≤ groupsize && (storage[local_idx] = val)
     @synchronize()
 
     s::UInt64 = groupsize ÷ 0x2
     while s > 0x0
         if (local_idx - 0x1) < s
             other_idx = local_idx + s
             if other_idx ≤ groupsize
-                storage[local_idx] = op(storage[local_idx], storage[other_idx])
+                @inbounds storage[local_idx] = op(storage[local_idx], storage[other_idx])
             end
         end
         @synchronize()
         s >>= 0x1
     end
 
     if local_idx == 0x1
-        val = storage[local_idx]
+        @inbounds val = storage[local_idx]
     end
     return val
 end
@@ -98,8 +92,9 @@ macro shfl_down(val, offset)
     end
 end
 
-# Backends should implement this.
+# Backends should implement these two.
 function __shfl_down end
+supports_warp_reduction(::CPU) = false
 
 @inline function __warp_reduce(val, op)
     offset::UInt32 = UInt32(32) ÷ 0x2
@@ -115,7 +110,7 @@ const __warpsize::UInt32 = 32
 # Maximum number of warps (for a groupsize = 1024).
 const __warp_bins::UInt32 = 32
 
-function __warp_groupreduce(__ctx__, op, val::T, neutral::T, ::Val{groupsize}) where {T, groupsize}
+function __groupreduce(__ctx__, op, val::T, neutral::T, ::Val{groupsize}, ::Val{:warp}) where {T, groupsize}
     storage = @localmem T __warp_bins
 
     local_idx = @index(Local)
@@ -124,12 +119,12 @@ function __warp_groupreduce(__ctx__, op, val::T, neutral::T, ::Val{groupsize}) w
 
     # Each warp performs a reduction and writes results into its own bin in `storage`.
     val = __warp_reduce(val, op)
-    lane == 0x1 && (storage[warp_id] = val)
+    @inbounds lane == 0x1 && (storage[warp_id] = val)
     @synchronize()
 
     # Final reduction of the `storage` on the first warp.
     within_storage = (local_idx - 0x1) < groupsize ÷ __warpsize
-    val = within_storage ? storage[lane] : neutral
+    @inbounds val = within_storage ? storage[lane] : neutral
     warp_id == 0x1 && (val = __warp_reduce(val, op))
     return val
 end

test/groupreduce.jl

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

@@ -41,57 +41,57 @@ 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
-    #     unittest_testsuite(backend, backend_str, backend_mod, DAT; skip_tests)
-    # end
+    @conditional_testset "Unittests" skip_tests begin
+        unittest_testsuite(backend, backend_str, backend_mod, DAT; skip_tests)
+    end
 
-    # @conditional_testset "SpecialFunctions" skip_tests begin
-    #     specialfunctions_testsuite(backend)
-    # end
+    @conditional_testset "SpecialFunctions" skip_tests begin
+        specialfunctions_testsuite(backend)
+    end
 
-    # @conditional_testset "Localmem" skip_tests begin
-    #     localmem_testsuite(backend, AT)
-    # end
+    @conditional_testset "Localmem" skip_tests begin
+        localmem_testsuite(backend, AT)
+    end
 
-    # @conditional_testset "Private" skip_tests begin
-    #     private_testsuite(backend, AT)
-    # end
+    @conditional_testset "Private" skip_tests begin
+        private_testsuite(backend, AT)
+    end
 
-    # @conditional_testset "Unroll" skip_tests begin
-    #     unroll_testsuite(backend, AT)
-    # end
+    @conditional_testset "Unroll" skip_tests begin
+        unroll_testsuite(backend, AT)
+    end
 
-    # @testset "NDIteration" begin
-    #     nditeration_testsuite()
-    # end
+    @testset "NDIteration" begin
+        nditeration_testsuite()
+    end
 
-    # @conditional_testset "copyto!" skip_tests begin
-    #     copyto_testsuite(backend, AT)
-    # end
+    @conditional_testset "copyto!" skip_tests begin
+        copyto_testsuite(backend, AT)
+    end
 
-    # @conditional_testset "Devices" skip_tests begin
-    #     devices_testsuite(backend)
-    # end
+    @conditional_testset "Devices" skip_tests begin
+        devices_testsuite(backend)
+    end
 
-    # @conditional_testset "Printing" skip_tests begin
-    #     printing_testsuite(backend)
-    # end
+    @conditional_testset "Printing" skip_tests begin
+        printing_testsuite(backend)
+    end
 
-    # @conditional_testset "Compiler" skip_tests begin
-    #     compiler_testsuite(backend, AT)
-    # end
+    @conditional_testset "Compiler" skip_tests begin
+        compiler_testsuite(backend, AT)
+    end
 
-    # @conditional_testset "Reflection" skip_tests begin
-    #     reflection_testsuite(backend, backend_str, AT)
-    # end
+    @conditional_testset "Reflection" skip_tests begin
+        reflection_testsuite(backend, backend_str, AT)
+    end
 
-    # @conditional_testset "Convert" skip_tests begin
-    #     convert_testsuite(backend, AT)
-    # end
+    @conditional_testset "Convert" skip_tests begin
+        convert_testsuite(backend, AT)
+    end
 
-    # @conditional_testset "Examples" skip_tests begin
-    #     examples_testsuite(backend_str)
-    # end
+    @conditional_testset "Examples" skip_tests begin
+        examples_testsuite(backend_str)
+    end
 
     @conditional_testset "@groupreduce" skip_tests begin
         groupreduce_testsuite(backend, AT)