DArray: Add accumulate, cumsum and cumprod

docs/src/darray.md

@@ -685,6 +685,7 @@ From `Base`:
 - `map`/`reduce`/`mapreduce`
 - `sum`/`prod`
 - `minimum`/`maximum`/`extrema`
+- `accumulate`/`cumsum`/`cumprod`
 - `map!`
 
 From `Random`:

src/array/map-reduce.jl

@@ -223,3 +223,62 @@ var(f::Function, x::DArray; dims=nothing) =
 std(x::DArray; dims=nothing) = std(identity, x; dims)
 std(f::Function, x::DArray; dims=nothing) =
     map(sqrt, var(f, x; dims))
+
+#### Accumulate
+
+# `Base`'s fallback scans element by element, which `getindex(::DArray, ...)`
+# cannot do. Scan each partition on its own instead, then combine the running
+# last slice of the partitions before it along `dims` into it. Like `reduce`,
+# this assumes `op` is associative.
+_accumulate_part(op, x, RT, dims, kw) = accumulate(op, _widen(x, RT); dims, kw...)
+# Only the first partition along `dims` starts from `init`. When `init` widens
+# the eltype, the others would scan in the narrower one and could overflow or
+# lose precision before the carry reaches them, so widen the partitions first.
+function _widen(x, ::Type{RT}) where RT
+    T = eltype(x)
+    widen = T <: Number && T !== RT && isconcretetype(RT) && promote_type(T, RT) === RT
+    return widen ? RT.(x) : x
+end
+_add_carry(op, carry, x) = op.(carry, x)
+function _next_carry(op, carry, x, dims)
+    n = size(x, dims)
+    slice = selectdim(x, dims, n:n)
+    return carry === nothing ? copy(slice) : op.(carry, slice)
+end
+
+function Base.accumulate(op, A::DArray; dims::Union{Nothing,Integer}=nothing, kw...)
+    if dims === nothing
+        # Without `dims`, `Base` scans any other array in linear order
+        ndims(A) == 1 || return invoke(accumulate, Tuple{Any,Any}, op, A; kw...)
+        return accumulate(op, A; dims=1, kw...)
+    end
+    dims > 0 || throw(ArgumentError("dims must be a positive integer"))
+    isempty(kw) || keys(kw) === (:init,) ||
+        throw(ArgumentError("accumulate does not support the keyword arguments $(setdiff(keys(kw), (:init,)))"))
+    return _accumulate(op, A, Int(dims), values(kw))
+end
+
+function _accumulate(op, A::DArray{T,N}, dims::Int, kw::NamedTuple) where {T,N}
+    # The eltype `Base` gives for a plain array
+    RT = eltype(Base.promote_op(x -> accumulate(op, x; dims, kw...), Array{T,N}))
+
+    parts = chunks(A)
+    thunks = similar(parts, DTask)
+    carries = similar(parts, DTask)
+    prev = CartesianIndex(ntuple(d -> Int(d == dims), N))
+    for I in CartesianIndices(parts)
+        k = dims <= N ? I[dims] : 1
+        scanned = Dagger.@spawn _accumulate_part(op, parts[I], RT, dims, k == 1 ? kw : (;))
+        carry = k == 1 ? nothing : carries[I - prev]
+        thunks[I] = k == 1 ? scanned : Dagger.@spawn _add_carry(op, carry, scanned)
+        if k < size(parts, dims)
+            carries[I] = Dagger.@spawn _next_carry(op, carry, scanned, dims)
+        end
+    end
+    return DArray(RT, domain(A), domainchunks(A), thunks, A.partitioning, A.concat)
+end
+
+Base.cumsum(A::DArray; dims::Integer) = accumulate(Base.add_sum, A; dims)
+Base.cumsum(A::DVector) = cumsum(A; dims=1)
+Base.cumprod(A::DArray; dims::Integer) = accumulate(Base.mul_prod, A; dims)
+Base.cumprod(A::DVector) = cumprod(A; dims=1)

test/array/mapreduce.jl

@@ -95,3 +95,36 @@ end
     @testset "prod" test_mapreduce_type((X; dims)->prod(x->x>0.05, X; dims), A; cmp=Base.:(==))
     @testset "maximum" test_mapreduce_type((X; dims)->maximum(x->x>0.95, X; dims), A; cmp=Base.:(==))
 end
+
+@testset "accumulate" begin
+    @testset "$(size(A)) in $bs" for (A, bs) in ((rand(1:9, 10), (3,)),
+                                                 (rand(1:9, 10, 9), (3, 4)),
+                                                 (rand(1:9, 5, 6, 7), (2, 3, 4)))
+        DA = distribute(A, Blocks(bs...))
+        @testset "dims=$dims" for dims in 1:ndims(A)+1
+            R = cumsum(DA; dims)
+            @test eltype(R) == eltype(cumsum(A; dims))
+            @test collect(R) == cumsum(A; dims)
+            @test collect(cumprod(DA; dims)) == cumprod(A; dims)
+            @test collect(accumulate(max, DA; dims)) == accumulate(max, A; dims)
+            R = accumulate(+, DA; dims, init=0.5)
+            @test eltype(R) == Float64
+            @test collect(R) == accumulate(+, A; dims, init=0.5)
+        end
+    end
+
+    v = rand(1:9, 10)
+    Dv = distribute(v, Blocks(3))
+    @test collect(cumsum(Dv)) == cumsum(v)
+    @test collect(cumprod(Dv)) == cumprod(v)
+    @test collect(accumulate(+, Dv; init=2)) == accumulate(+, v; init=2)
+    v8 = rand(Int8, 10)
+    @test eltype(cumsum(distribute(v8, Blocks(3)))) == Int
+    @test collect(cumsum(distribute(v8, Blocks(3)))) == cumsum(v8)
+    # An `init` of a wider type widens every partition, not just the first
+    w8 = fill(Int8(100), 10, 4)
+    @test collect(accumulate(+, distribute(w8, Blocks(3, 2)); dims=1, init=0)) == accumulate(+, w8; dims=1, init=0)
+    @test collect(accumulate(+, distribute(w8[:, 1], Blocks(3)); init=0)) == accumulate(+, w8[:, 1]; init=0)
+    @test_throws ArgumentError cumsum(Dv; dims=0)
+    @test_throws ArgumentError accumulate(+, Dv; foo=1)
+end

test/gpu.jl

@@ -36,6 +36,28 @@ end
     end
 end
 
+@testset "DArray accumulate ($name)" for (name, scope) in GPU_SCOPES
+    Dagger.with_options(; scope) do
+        DA = rand(Blocks(4, 4), Float32, 16, 16)
+        A = collect(DA)
+        if name == :OpenCL
+            # OpenCL.jl has no `accumulate` for `CLArray`
+            @test_broken collect(cumsum(DA; dims=1)) ≈ cumsum(A; dims=1)
+            return
+        end
+        @test collect(cumsum(DA; dims=1)) ≈ cumsum(A; dims=1)
+        @test collect(cumprod(DA; dims=2)) ≈ cumprod(A; dims=2)
+        @test collect(accumulate(max, DA; dims=1)) == accumulate(max, A; dims=1)
+        @test collect(accumulate(+, DA; dims=2, init=1f0)) ≈ accumulate(+, A; dims=2, init=1f0)
+        # Metal has no Float64
+        if name != :Metal
+            @test collect(accumulate(+, DA; dims=1, init=0.0)) ≈ accumulate(+, A; dims=1, init=0.0)
+        end
+        DV = rand(Blocks(4), Float32, 16)
+        @test collect(cumsum(DV)) ≈ cumsum(collect(DV))
+    end
+end
+
 @everywhere begin
     function isongpu(X)
         return !(X isa Array)