From 9951e88cda11764ac7db284d44b9f7842d2ae71b Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Tue, 13 Oct 2026 11:11:24 +0300 Subject: [PATCH 1/2] DArray: Add accumulate, cumsum and cumprod There was no `accumulate` method for a `DArray`, so `Base`'s generic one scanned it element by element. With `dims` that throws a `ConcurrencyViolationError`, and on a vector it returns a `DVector{Any}`. Scan each partition on its own, then combine the running last slice of the partitions before it along `dims` into it. Only the carries run in sequence, and each one is a single slice. Like `reduce`, this assumes `op` is associative. Only the first partition starts from `init`, so when `init` widens the eltype the partitions are widened before the scan. --- docs/src/darray.md | 1 + src/array/map-reduce.jl | 59 +++++++++++++++++++++++++++++++++++++++++ test/array/mapreduce.jl | 33 +++++++++++++++++++++++ test/gpu.jl | 14 ++++++++++ 4 files changed, 107 insertions(+) diff --git a/docs/src/darray.md b/docs/src/darray.md index bba33a17b..76d909290 100644 --- a/docs/src/darray.md +++ b/docs/src/darray.md @@ -685,6 +685,7 @@ From `Base`: - `map`/`reduce`/`mapreduce` - `sum`/`prod` - `minimum`/`maximum`/`extrema` +- `accumulate`/`cumsum`/`cumprod` - `map!` From `Random`: diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 739a7a448..9555fca1e 100644 --- a/src/array/map-reduce.jl +++ b/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) diff --git a/test/array/mapreduce.jl b/test/array/mapreduce.jl index e16263ce1..64332ed39 100644 --- a/test/array/mapreduce.jl +++ b/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 diff --git a/test/gpu.jl b/test/gpu.jl index b6f236c7f..83709b276 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -36,6 +36,20 @@ 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) + @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) + @test collect(accumulate(+, DA; dims=1, init=0.0)) ≈ accumulate(+, A; dims=1, init=0.0) + DV = rand(Blocks(4), Float32, 16) + @test collect(cumsum(DV)) ≈ cumsum(collect(DV)) + end +end + @everywhere begin function isongpu(X) return !(X isa Array) From b0fd3502e7707647ed8e97b2a95403e90f9dabcd Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Tue, 13 Oct 2026 15:13:15 +0300 Subject: [PATCH 2/2] Mark accumulate broken on OpenCL, skip the Float64 test on Metal OpenCL.jl has no `accumulate` for `CLArray`, so `Base` falls back to scalar indexing. Metal has no Float64, so the `init=0.0` test can't run there. --- test/gpu.jl | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/test/gpu.jl b/test/gpu.jl index 83709b276..f1802b328 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -40,11 +40,19 @@ end 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) - @test collect(accumulate(+, DA; dims=1, init=0.0)) ≈ accumulate(+, A; dims=1, init=0.0) + # 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