Faster cumsum & cumprod. Fixes #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

@@ -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.