Merge pull request JuliaLang#7359 from JuliaLang/teh/cumsumprod

Faster cumsum & cumprod. Fixes JuliaLang#7342.

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
 

base/exports.jl

@@ -506,7 +506,9 @@ export
  cummax,
  cummin,
  cumprod,
+ cumprod!,
  cumsum,
+ cumsum!,
  cumsum_kbn,
  extrema,
  fill!,

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
 

doc/stdlib/base.rst

@@ -4040,10 +4040,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.