Skip to content

Commit 9b8b455

Browse files
committed
FIX: Broken debug assertion in collect to f-order array
This fixes a "bug" in layout_impl - if an array is effectively one-dimensional, like for example has dimension 2 x 1 x 1 (3d array with len two), then it's possible it is both C and F-contig. This was the reason for the debug assertion firing: collection starts into an f-order output array, but when we've split the array enough times, it's still an f-order array but in shape 1 x 1 x 1 or 2 x 1 x 1 in this case - and we detected that as a c-order array.
1 parent 2af780f commit 9b8b455

File tree

2 files changed

+30
-12
lines changed

2 files changed

+30
-12
lines changed

src/zip/mod.rs

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ where
5858
pub(crate) fn layout_impl(&self) -> Layout {
5959
let n = self.ndim();
6060
if self.is_standard_layout() {
61-
if n <= 1 {
61+
// effectively one-dimensional => C and F layout compatible
62+
if n <= 1 || self.shape().iter().filter(|&&len| len > 1).count() <= 1 {
6263
Layout::one_dimensional()
6364
} else {
6465
Layout::c()
@@ -1173,8 +1174,19 @@ macro_rules! map_impl {
11731174
{
11741175
// Get the last producer; and make a Partial that aliases its data pointer
11751176
let (.., ref output) = &self.parts;
1176-
debug_assert!(output.layout().is(CORDER | FORDER));
1177-
debug_assert_eq!(output.layout().tendency() >= 0, self.layout_tendency >= 0);
1177+
1178+
// debug assert that the output is contiguous in the memory layout we need
1179+
if cfg!(debug_assertions) {
1180+
let out_layout = output.layout();
1181+
assert!(out_layout.is(CORDER | FORDER));
1182+
assert!(
1183+
(self.layout_tendency <= 0 && out_layout.tendency() <= 0) ||
1184+
(self.layout_tendency >= 0 && out_layout.tendency() >= 0),
1185+
"layout tendency violation for self layout {:?}, output layout {:?},\
1186+
output shape {:?}",
1187+
self.layout, out_layout, output.raw_dim());
1188+
}
1189+
11781190
let mut partial = Partial::new(output.as_ptr());
11791191

11801192
// Apply the mapping function on this zip

tests/par_zip.rs

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -107,20 +107,26 @@ fn test_zip_small_collect() {
107107

108108
for m in 0..32 {
109109
for n in 0..16 {
110-
let dim = (m, n);
111-
let b = Array::from_shape_fn(dim, |(i, j)| 1. / (i + 2 * j + 1) as f32);
112-
let c = Array::from_shape_fn(dim, |(i, j)| f32::ln((1 + i + j) as f32));
113-
114-
{
115-
let a = Zip::from(&b).and(&c).par_map_collect(|x, y| x + y);
116-
117-
assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
118-
assert_eq!(a.strides(), b.strides());
110+
for &is_f in &[false, true] {
111+
let dim = (m, n).set_f(is_f);
112+
let b = Array::from_shape_fn(dim, |(i, j)| 1. / (i + 2 * j + 1) as f32);
113+
let c = Array::from_shape_fn(dim, |(i, j)| f32::ln((1 + i + j) as f32));
114+
115+
{
116+
let a = Zip::from(&b).and(&c).par_map_collect(|x, y| x + y);
117+
118+
assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
119+
if m > 1 && n > 1 {
120+
assert_eq!(a.strides(), b.strides(),
121+
"Failure for {}x{} c/f: {:?}", m, n, is_f);
122+
}
123+
}
119124
}
120125
}
121126
}
122127
}
123128

129+
124130
#[test]
125131
#[cfg(feature = "approx")]
126132
fn test_zip_assign_into() {

0 commit comments

Comments
 (0)