Port over accumulations

Also tests

Author: christiangnrd

src/GPUArrays.jl

@@ -28,6 +28,7 @@ include("host/construction.jl")
 include("host/base.jl")
 include("host/indexing.jl")
 include("host/broadcast.jl")
+include("host/accumulate.jl")
 include("host/mapreduce.jl")
 include("host/linalg.jl")
 include("host/math.jl")

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::AnyGPUVector, 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::AnyGPUArray;
+                         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

@@ -90,6 +90,7 @@ include("testsuite/indexing.jl")
 include("testsuite/base.jl")
 include("testsuite/vector.jl")
 include("testsuite/reductions.jl")
+include("testsuite/accumulations.jl")
 include("testsuite/broadcasting.jl")
 include("testsuite/linalg.jl")
 include("testsuite/math.jl")

test/testsuite/accumulations.jl

@@ -0,0 +1,108 @@
+@testsuite "accumulations" (AT, eltypes)->begin
+    @testset "$ET" for ET in eltypes
+        range = ET <: Real ? (ET(1):ET(10)) : ET
+
+        # 1d arrays
+        for num_elems in 1:256
+            @test compare(A->accumulate(+, A; init=zero(ET)), AT, rand(range, num_elems))
+        end
+
+        for num_elems = rand(1:100, 10)
+            @test compare(A->accumulate(+, A; init=zero(ET)), AT, rand(range, num_elems))
+        end
+
+        for _ in 1:10 # nd arrays reduced as 1d
+            n1 = rand(1:10)
+            n2 = rand(1:10)
+            n3 = rand(1:10)
+            @test compare(A->accumulate(+, A; init=zero(ET)), AT, rand(range, n1, n2, n3))
+        end
+
+        for num_elems = rand(1:100, 10) # init value
+            init = rand(range)
+            @test compare(A->accumulate(+, A; init), AT, rand(range, num_elems))
+        end
+
+
+        # nd arrays
+        for dims in 1:4 # corner cases
+            for isize in 1:3
+                for jsize in 1:3
+                    for ksize in 1:3
+                        @test compare(A->accumulate(+, A; dims, init=zero(ET)), AT, rand(range, isize, jsize, ksize))
+                    end
+                end
+            end
+        end
+
+        for _ in 1:10
+            for dims in 1:3
+                n1 = rand(1:10)
+                n2 = rand(1:10)
+                n3 = rand(1:10)
+                @test compare(A->accumulate(+, A; dims, init=zero(ET)), AT, rand(range, n1, n2, n3))
+            end
+        end
+
+        for _ in 1:10 # init value
+            for dims in 1:3
+                n1 = rand(1:10)
+                n2 = rand(1:10)
+                n3 = rand(1:10)
+                init = rand(range)
+                @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
+            # Skip large tests on small datatypes
+            n >= 10000 && sizeof(real(ET)) <= 2 && continue
+
+            @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
+
+        # in place
+        @test compare(x->(accumulate!(+, x, copy(x)); x), AT, rand(range, 2))
+
+        @test_throws ArgumentError("accumulate does not support the keyword arguments [:bad_kwarg]") accumulate(+, AT(rand(ET, 10)); bad_kwarg="bad")
+    end
+end
+
+@testsuite "accumulations/cumsum & cumprod" (AT, eltypes)->begin
+    @test compare(cumsum, AT, rand(Bool, 16))
+
+    @testset "$ET" for ET in eltypes
+        range = ET <: Real ? (ET(1):ET(10)) : ET
+
+        # cumsum
+        for num_elems in rand(1:100, 10)
+            @test compare(A->cumsum(A; dims=1), AT, rand(range, num_elems))
+        end
+
+        for _ in 1:10
+            for dims in 1:3
+                n1 = rand(1:10)
+                n2 = rand(1:10)
+                n3 = rand(1:10)
+                @test compare(A->cumsum(A; dims), AT, rand(range, n1, n2, n3))
+            end
+        end
+
+
+        # cumprod
+        range = ET <: Real ? (ET(1):ET(10)) : ET
+        @test compare(A->cumprod(A; dims=1), AT, ones(ET, 100_000))
+
+        for _ in 1:10
+            for dims in 1:3
+                n1 = rand(1:10)
+                n2 = rand(1:10)
+                n3 = rand(1:10)
+                @test compare(A->cumprod(A; dims), AT, rand(range, n1, n2, n3))
+            end
+        end
+    end
+end