Port over accumulations

Author: christiangnrd

src/host/accumulate.jl

@@ -0,0 +1,35 @@
+## Base interface
+
+Base._accumulate!(op, output::AnyGPUArray, input::AnyGPUVector, dims::Nothing, init::Nothing) =
+    AK.accumulate!(op, output, input, get_backend(output); dims, init=AK.neutral_element(op, eltype(output)))
+
+Base._accumulate!(op, output::AnyGPUArray, input::AnyGPUArray, dims::Integer, init::Nothing) =
+    AK.accumulate!(op, output, input, get_backend(output); dims, init=AK.neutral_element(op, eltype(output)))
+
+Base._accumulate!(op, output::AnyGPUArray, input::MtlVector, dims::Nothing, init::Some) =
+    AK.accumulate!(op, output, input, get_backend(output); dims, init=something(init))
+
+Base._accumulate!(op, output::AnyGPUArray, input::AnyGPUArray, dims::Integer, init::Some) =
+    AK.accumulate!(op, output, input, get_backend(output); dims, init=something(init))
+
+Base.accumulate_pairwise!(op, result::AnyGPUVector, v::AnyGPUVector) = accumulate!(op, result, v)
+
+# default behavior unless dims are specified by the user
+function Base.accumulate(op, A::WrappedGPUArray;
+                         dims::Union{Nothing,Integer}=nothing, kw...)
+    nt = values(kw)
+    if dims === nothing && !(A isa AbstractVector)
+        # This branch takes care of the cases not handled by `_accumulate!`.
+        return reshape(AK.accumulate(op, A[:], get_backend(A); init = (:init in keys(kw) ? nt.init : AK.neutral_element(op, eltype(A)))), size(A))
+    end
+    if isempty(kw)
+        out = similar(A, Base.promote_op(op, eltype(A), eltype(A)))
+        init = AK.neutral_element(op, eltype(out))
+    elseif keys(nt) === (:init,)
+        out = similar(A, Base.promote_op(op, typeof(nt.init), eltype(A)))
+        init = nt.init
+    else
+        throw(ArgumentError("accumulate does not support the keyword arguments $(setdiff(keys(nt), (:init,)))"))
+    end
+    AK.accumulate!(op, out, A, get_backend(A); dims, init)
+end

test/testsuite.jl

@@ -96,6 +96,7 @@ include("testsuite/random.jl")
 include("testsuite/uniformscaling.jl")
 include("testsuite/statistics.jl")
 include("testsuite/alloc_cache.jl")
+include("testsuite/accumulations.jl")
 include("testsuite/jld2ext.jl")
 
 """

test/testsuite/accumulations.jl

@@ -1,4 +1,3 @@
-# @testsuite "accumulations" (AT, eltypes)->begin
 @testsuite "accumulations" (AT, eltypes)->begin
     @testset "$ET" for ET in eltypes
         range = ET <: Real ? (ET(1):ET(10)) : ET
@@ -54,6 +53,13 @@
                 @test compare(A->accumulate(+, A; init, dims), AT, rand(range, n1, n2, n3))
             end
         end
+
+        # Larger containers to try and detect weird bugs
+        for n in (0, 1, 2, 3, 10, 10_000, 16384, 16384+1) # small, large, odd & even, pow2 and not
+            @test compare(x->accumulate(+, x), AT, rand(range, n))
+            @test compare(x->accumulate(+, x), AT, rand(range, n, 2))
+            @test compare(Base.Fix2((x,y)->accumulate(+, x; init=y), rand(range)), AT, rand(range, n))
+        end
     end
 end