Skip to content

Commit d946360

Browse files
committed
append: performance optimize stride check
This .max_by_key() call showed up in profiling, and this change improves the select_axis0 benchmark by reducing runtime by 15%
1 parent 1206bd6 commit d946360

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

src/impl_owned_array.rs

+12-5
Original file line numberDiff line numberDiff line change
@@ -438,12 +438,19 @@ impl<A, D> Array<A, D>
438438
// array must be empty or have `axis` as the outermost (longest stride) axis
439439
if !self_is_empty && current_axis_len > 1 {
440440
// `axis` must be max stride axis or equal to its stride
441-
let max_axis = self.axes().max_by_key(|ax| ax.stride.abs()).unwrap();
442-
if max_axis.axis != axis && max_axis.stride.abs() > self.stride_of(axis) {
443-
incompatible_layout = true;
444-
}
445-
if self.stride_of(axis) < 0 {
441+
let axis_stride = self.stride_of(axis);
442+
if axis_stride < 0 {
446443
incompatible_layout = true;
444+
} else {
445+
for ax in self.axes() {
446+
if ax.axis == axis {
447+
continue;
448+
}
449+
if ax.len > 1 && ax.stride.abs() > axis_stride {
450+
incompatible_layout = true;
451+
break;
452+
}
453+
}
447454
}
448455
}
449456

0 commit comments

Comments
 (0)