diff --git a/src/host/accumulate.jl b/src/host/accumulate.jl index 9212699f..1c349943 100644 --- a/src/host/accumulate.jl +++ b/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 diff --git a/src/host/broadcast.jl b/src/host/broadcast.jl index 7132863d..4498c6e6 100644 --- a/src/host/broadcast.jl +++ b/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 diff --git a/test/testsuite/accumulate.jl b/test/testsuite/accumulate.jl index 70a567aa..92f3662e 100644 --- a/test/testsuite/accumulate.jl +++ b/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 diff --git a/test/testsuite/broadcasting.jl b/test/testsuite/broadcasting.jl index e892c69f..19e8b192 100644 --- a/test/testsuite/broadcasting.jl +++ b/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)