sum_axis very slow when summing on last axis

[gaumut]

On my computer

let a = Array3::<f64>::zeros((500,600,700));
let b = a.sum_axis(Axis(2));

is about 9 times slower (500ms vs 55ms) than

let a = Array3::<f64>::zeros((500,600,700));
let mut c = Array2::<f64>::zeros((500,600));
Zip::from(&mut c).and(a.lanes(Axis(2))).for_each(|c,a| {
    *c = a.sum();
});

On the other hand when summing on Axis(0) sum_axis is about two times faster than Zip.

The gain seems big enough to me to warrant an optimization but I'm not sure how to deal with intermediate cases (summing on some middle axis). Or can we just add the optimization for the summation on the last axis as it's a very common case ?

(Of course it's the other way around for a column major array.)

[jturner314]

I've created #1061 in an attempt to improve on this. The heuristics could be improved further, but it should be better than the current implementation. Could you do some benchmarking to confirm that it's an improvement? You can test the PR branch by adding the following to your Cargo.toml:

[patch.crates-io]
ndarray = { git = "https://github.com/jturner314/ndarray.git", branch = "faster-sum_axis" }

[gaumut]

Ok this is a quick and dirty benchmark with the linux time command and with this code

use ndarray::{Array3,Axis,Zip,Array2, ShapeBuilder};
fn main() {
    let a = Array3::<f64>::zeros((700,700,700).f());
    let b = a.sum_axis(Axis(2));
    /*  
    let mut b = Array2::<f64>::zeros((700,700));
    Zip::from(&mut b).and(a.lanes(Axis(2))).for_each(|b, a| {
        *b = a.sum();
    });
    */
}

Times in ms :

ROW_MAJOR	Axis(0)	Axis(1)	Axis(2)
sum_axis	165	175	550
sum_axis_patched	161	177	76
Zip	398	385	74

COLUMN_MAJOR	Axis(0)	Axis(1)	Axis(2)
sum_axis	586	367	282
sum_axis_patched	78	372	283
Zip	77	415	451

Seems good to me, thanks.