From bd7d73f6b11bacff194f6d2b121c1df55bb2ab88 Mon Sep 17 00:00:00 2001
From: Mathieu Poumeyrol <kali@zoy.org>
Date: Fri, 16 Feb 2024 09:35:44 +0100
Subject: [PATCH] fix packing

---
 core/src/ops/cnn/conv/conv.rs        |  3 +-
 core/src/ops/cnn/conv/lazy_im2col.rs | 43 ++++++----------------------
 test-rt/suite-unit/src/conv_f32.rs   | 18 ++++++++++++
 3 files changed, 29 insertions(+), 35 deletions(-)

diff --git a/core/src/ops/cnn/conv/conv.rs b/core/src/ops/cnn/conv/conv.rs
index c5c68d21e1..28a5b75049 100644
--- a/core/src/ops/cnn/conv/conv.rs
+++ b/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())?;
diff --git a/core/src/ops/cnn/conv/lazy_im2col.rs b/core/src/ops/cnn/conv/lazy_im2col.rs
index 1a565de0bd..7ff9dc29c9 100644
--- a/core/src/ops/cnn/conv/lazy_im2col.rs
+++ b/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,
-    )
-    }
-    }
-    */
 }
diff --git a/test-rt/suite-unit/src/conv_f32.rs b/test-rt/suite-unit/src/conv_f32.rs
index a33989e5f2..1e7d43bde0 100644
--- a/test-rt/suite-unit/src/conv_f32.rs
+++ b/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)
 }