diff --git a/base/array.jl b/base/array.jl index 9d8bf7dbe3b48..4ec7dbbec2e58 100644 --- a/base/array.jl +++ b/base/array.jl @@ -1408,33 +1408,6 @@ for (f, fp, op) = ((:cumsum, :cumsum_pairwise, :+), return c end - @eval function ($f)(A::StridedArray, axis::Integer) - dimsA = size(A) - ndimsA = ndims(A) - axis_size = dimsA[axis] - axis_stride = 1 - for i = 1:(axis-1) - axis_stride *= size(A,i) - end - - B = $(op===:+ ? (:(similar(A,_cumsum_type(A)))) : - (:(similar(A)))) - - if axis_size < 1 - return B - end - - for i = 1:length(A) - if div(i-1, axis_stride) % axis_size == 0 - B[i] = A[i] - else - B[i] = ($op)(B[i-axis_stride], A[i]) - end - end - - return B - end - @eval ($f)(A::AbstractArray) = ($f)(A, 1) end diff --git a/base/exports.jl b/base/exports.jl index de94b499e1430..816aaf07c6704 100644 --- a/base/exports.jl +++ b/base/exports.jl @@ -506,7 +506,9 @@ export cummax, cummin, cumprod, + cumprod!, cumsum, + cumsum!, cumsum_kbn, extrema, fill!, diff --git a/base/multidimensional.jl b/base/multidimensional.jl index cde16603f0e9b..c25e3984a31fe 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -134,6 +134,40 @@ end eval(ngenerate(:N, nothing, :(setindex!{T}(s::SubArray{T,N}, v, ind::Integer)), gen_setindex!_body, 2:5, false)) + cumsum(A::AbstractArray, axis::Integer) = cumsum!(similar(A, Base._cumsum_type(A)), A, axis) +cumprod(A::AbstractArray, axis::Integer) = cumprod!(similar(A), A, axis) + +for (f, op) in ((:cumsum!, :+), + (:cumprod!, :*)) + @eval begin + @ngenerate N typeof(B) function ($f){T,N}(B, A::AbstractArray{T,N}, axis::Integer) + if size(B, axis) < 1 + return B + end + size(B) == size(A) || throw(DimensionMismatch("Size of B must match A")) + if axis == 1 + # We can accumulate to a temporary variable, which allows register usage and will be slightly faster + @inbounds @nloops N i d->(d > 1 ? (1:size(A,d)) : (1:1)) begin + tmp = convert(eltype(B), @nref(N, A, i)) + @nref(N, B, i) = tmp + for i_1 = 2:size(A,1) + tmp = ($op)(tmp, @nref(N, A, i)) + @nref(N, B, i) = tmp + end + end + else + @nexprs N d->(isaxis_d = axis == d) + # Copy the initial element in each 1d vector along dimension `axis` + @inbounds @nloops N i d->(d == axis ? (1:1) : (1:size(A,d))) @nref(N, B, i) = @nref(N, A, i) + # Accumulate + @inbounds @nloops N i d->((1+isaxis_d):size(A, d)) d->(j_d = i_d - isaxis_d) begin + @nref(N, B, i) = ($op)(@nref(N, B, j), @nref(N, A, i)) + end + end + B + end + end +end ### from abstractarray.jl diff --git a/doc/stdlib/base.rst b/doc/stdlib/base.rst index d79562a86eabf..aa12e6e4f729c 100644 --- a/doc/stdlib/base.rst +++ b/doc/stdlib/base.rst @@ -4038,10 +4038,18 @@ Array functions Cumulative product along a dimension. +.. function:: cumprod!(B, A, [dim]) + + Cumulative product of ``A`` along a dimension, storing the result in ``B``. + .. function:: cumsum(A, [dim]) Cumulative sum along a dimension. +.. function:: cumsum!(B, A, [dim]) + + Cumulative sum of ``A`` along a dimension, storing the result in ``B``. + .. function:: cumsum_kbn(A, [dim]) Cumulative sum along a dimension, using the Kahan-Babuska-Neumaier compensated summation algorithm for additional accuracy.