Follow Base when accumulating along a dimension beyond the array's

src/host/accumulate.jl

@@ -3,9 +3,12 @@
 # Base's `accumulate!` and `accumulate` reach these four methods, split as Base's are. `init` is
 # passed on only when given (`Some(x)`), so that an explicit `init=nothing` stays an initial
 # value. Base's rules are kept here: without `dims`, only vectors are scanned, and a `dims` beyond
-# the array's copies it, ignoring `init` (AcceleratedKernels would apply `init` to every element).
-# `B`'s element type, which Base's front-end chose, sets the running values' type in
-# AcceleratedKernels.
+# the array's is a trailing singleton dimension, so every element starts a scan of its own: it is
+# combined with `init`, or passed through `reduce_first` (Julia versions without
+# JuliaLang/julia#64233 copy the array instead, ignoring `init`). `B`'s element type, which Base's
+# front-end chose, sets the running values' type in AcceleratedKernels.
+
+const ACCUMULATE_BEYOND_NDIMS_USES_INIT = accumulate(+, [0]; dims=2, init=1) == [1]
 
 Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUVector, dims::Nothing, init::Nothing) =
     AK.accumulate!(op, B, A)
@@ -14,13 +17,20 @@ Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUVector, dims::Nothing, init::Some
 function Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUArray, dims::Integer, init::Nothing)
     dims > 0 || throw(ArgumentError("dims must be a positive integer"))
     axes(B) == axes(A) || throw(DimensionMismatch("shape of B must match A"))
-    dims > ndims(A) && return copyto!(B, A)
+    if dims > ndims(A)
+        ACCUMULATE_BEYOND_NDIMS_USES_INIT || return copyto!(B, A)
+        return map!(x -> Base.reduce_first(op, x), B, A)
+    end
     AK.accumulate!(op, B, A; dims)
 end
 function Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUArray, dims::Integer, init::Some)
     dims > 0 || throw(ArgumentError("dims must be a positive integer"))
     axes(B) == axes(A) || throw(DimensionMismatch("shape of B must match A"))
-    dims > ndims(A) && return copyto!(B, A)
+    if dims > ndims(A)
+        ACCUMULATE_BEYOND_NDIMS_USES_INIT || return copyto!(B, A)
+        v = something(init)
+        return map!(x -> op(v, x), B, A)
+    end
     AK.accumulate!(op, B, A; dims, init=something(init))
 end
 

src/host/broadcast.jl

@@ -114,7 +114,7 @@ function Base.map!(f, dest::AnyGPUArray, xs::AbstractArray...)
     # (avoids the reshape + view that our mapreduce impl has to do)
     indices = LinearIndices.((dest, xs...))
     common_length = minimum(length.(indices))
-    common_length==0 && return
+    common_length==0 && return dest
 
     bc = Broadcast.instantiate(Broadcast.broadcasted(f, xs...))
     if bc isa Broadcast.Broadcasted

test/testsuite/accumulate.jl

@@ -41,8 +41,8 @@
     @test compare(A -> accumulate(something_add, A; init=nothing), AT, rand(1:10, 100))
 
     # Base's rules, which AcceleratedKernels leaves to GPUArrays: element types (which differ
-    # between Julia versions), a `dims` beyond the array's (which copies, ignoring `init`), and
-    # the running type Base's front-end chooses through the destination
+    # between Julia versions), a `dims` beyond the array's (where JuliaLang/julia#64233 started
+    # applying `init`), and the running type Base's front-end chooses through the destination
     for (f, x) in ((A -> accumulate(+, A), Int8[1, 2, 100]), (A -> accumulate(+, A; init=0), Int8[1, 2, 100]),
                    (cumsum, Int8[1, 2, 100]), (cumsum, Bool[1, 1, 0]), (cumprod, Int8[2, 3]),
                    (A -> accumulate(*, A), UInt8[16, 16, 16]),
@@ -57,6 +57,7 @@
                         zeros(Int32, 1, 3), Float32[0.5 1.0 2.0])
     @test compare_exact((B, A) -> accumulate!(+, B, A; dims=3, init=10), AT, zeros(Int, 3, 4),
                         rand(1:10, 3, 4))
+    @test compare_exact(A -> accumulate(+, A; dims=3, init=10), AT, zeros(Int, 0, 4))
 
     # Associative operators need not be commutative
     @test compare(A -> accumulate((a, b) -> a, A), AT, rand(1:10, 1000))
@@ -71,4 +72,5 @@
     end
     @test compare(A -> accumulate(compose, A; dims=2, init=(1, 0)), AT, affine(7, 10_000))
     @test compare(A -> accumulate(compose, A; dims=1), AT, affine(10_000, 7))
+    @test compare(A -> accumulate(compose, A; dims=3, init=(2, 1)), AT, affine(7, 10))
 end

test/testsuite/broadcasting.jl

@@ -166,6 +166,9 @@ function broadcasting(AT, eltypes)
             @test compare(AT, rand(ET, 2), 1:2) do x, y
                 map!(+, x, y)
             end
+            @test compare(AT, rand(ET, 0), rand(ET, 0)) do x, y
+                map!(+, x, y)
+            end
         end
 
         @testset "map $ET" begin
@@ -178,6 +181,9 @@ function broadcasting(AT, eltypes)
             @test compare(AT, rand(ET, 2,2), rand(ET, 2)) do x,y
                 map(+, x, y)
             end
+            @test compare(AT, rand(ET, 0), rand(ET, 2)) do x, y
+                map(+, x, y)
+            end
             # issue #580
             @test compare(AT, rand(ET, 2), 1:2) do x, y
                 map(+, x, y)