fix packing

core/src/ops/cnn/conv/conv.rs

@@ -362,7 +362,8 @@ impl Conv {
         let &[mut x, kernel, bias] = wire else { bail!("Wrong number of inputs") };
         let mut x_fact = model.outlet_fact(x)?.clone();
         let k_fact = model.outlet_fact(kernel)?.clone();
-        let (geo, _, k, _, mmm) = self.compute_geo(&k_fact, &x_fact)?;
+        let (geo, m, k, n, mmm) = self.compute_geo(&k_fact, &x_fact)?;
+        debug!("{name} as lazy_im2col: m={m} k={k} n={n} {mmm}");
         let input_shape = x_fact.shape.as_concrete().unwrap().to_vec();
         let mut geo = geo.to_concrete(&input_shape)?.into_owned();
         let mut input_shape: DataShape = self.pool_spec.data_format.shape(input_shape.into())?;

core/src/ops/cnn/conv/lazy_im2col.rs

@@ -242,44 +242,19 @@ impl<T: Datum + Copy> InputStore for LazyIm2col<T> {
     }
 
     fn panel(&self, i: usize, buffer: Option<*mut u8>) -> *const u8 {
+        let mn_start = i * self.packer.r;
+        let mn_end = (mn_start + self.packer.r).min(self.n);
+        let mn_range = mn_start as isize..mn_end as isize;
+        let k_range = 0..self.k as isize;
         let packed = buffer.unwrap();
-        let mn_end = ((i + 1) * self.packer.r).min(self.n) as isize;
-        let mn_range = (i * self.packer.r) as isize..mn_end;
-        if mn_range.len() == self.packer.r && mn_range.start % self.packer.r as isize == 0 {
+        if mn_range.len() == self.packer.r && mn_start % self.packer.r == 0 {
             let mut writer = self.packer.write_single_panel_with_k_outer(packed as *mut T);
-            self.write(&mut writer, 0..self.k as isize, mn_range)
+            self.write(&mut writer, k_range, mn_range);
         } else {
-            let mut writer = self.packer.write_with_k_outer(packed as *mut T, self.k, self.n);
-            self.write(&mut writer, 0..self.k as isize, mn_range)
+            let mut writer =
+                self.packer.write_with_k_outer(packed as *mut T, k_range.len(), mn_range.len());
+            self.write(&mut writer, k_range, mn_range);
         }
         packed
     }
-    /*
-    fn input(
-    &self,
-    packer: &tract_linalg::frame::Packer,
-    packed: *mut u8,
-    k_range: std::ops::Range<usize>,
-    mn_range: std::ops::Range<usize>,
-    ) {
-    let mn_end = mn_range.end.min(self.n) as isize;
-    let n_range = mn_range.start as isize..mn_end;
-    if n_range.len() == packer.r && mn_range.start % packer.r == 0 {
-    let mut writer = packer.write_single_panel_with_k_outer(packed as *mut T);
-    self.write(
-    &mut writer,
-    k_range.start as isize..k_range.end as isize,
-    mn_range.start as isize..n_range.end,
-    )
-    } else {
-    let mut writer =
-    packer.write_with_k_outer(packed as *mut T, k_range.len(), n_range.len());
-    self.write(
-    &mut writer,
-    k_range.start as isize..k_range.end as isize,
-    mn_range.start as isize..n_range.end,
-    )
-    }
-    }
-    */
 }

test-rt/suite-unit/src/conv_f32.rs

@@ -1196,5 +1196,23 @@ pub fn suite() -> TractResult<TestSuite> {
         },
     );
 
+    let mut data = Tensor::zero::<f32>(&[1, 5, 6]).unwrap();
+    *data.as_slice_mut::<f32>().unwrap().last_mut().unwrap() = 1.0;
+    let mut kernel = Tensor::zero::<f32>(&[1, 1, 3, 2]).unwrap();
+    *kernel.as_slice_mut::<f32>().unwrap().last_mut().unwrap() = 1.0;
+    suite.add(
+        "pack_0",
+        ConvProblem {
+            shape_in: DataFormat::CHW.from_n_c_hw(1, 1, [5, 6]).unwrap(),
+            kernel_format: KernelFormat::OIHW,
+            group: 1,
+            data: data.into_array::<f32>().unwrap(),
+            kernel: kernel.into_array::<f32>().unwrap(),
+            bias: None,
+            pad: PaddingSpec::Valid,
+            strides: tvec!(1, 1),
+        },
+    );
+
     Ok(suite)
 }