From 7e8193869ad1038a10b009d728f621bf1d3689ae Mon Sep 17 00:00:00 2001
From: Mathieu Poumeyrol <kali@zoy.org>
Date: Tue, 30 Jan 2024 13:03:16 +0100
Subject: [PATCH] stronger alignment for packing

---
 linalg/src/frame/pack.rs         |  6 +++---
 test-rt/suite-unit/src/conv_q.rs | 15 +++++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/linalg/src/frame/pack.rs b/linalg/src/frame/pack.rs
index e14dfcac8d..c5c64c80f7 100644
--- a/linalg/src/frame/pack.rs
+++ b/linalg/src/frame/pack.rs
@@ -24,11 +24,11 @@ impl Packer {
     }
 
     pub fn len<D: DimLike>(&self, k: D, n: D) -> D {
-        (n.divceil(self.r) * (k + self.end_padding_record)) * self.r
+        n.divceil(self.r) * self.single_panel_len(k)
     }
 
-    pub fn single_panel_len(&self, k: usize) -> usize {
-        (k + self.end_padding_record) * self.r
+    pub fn single_panel_len<D: DimLike>(&self, k: D) -> D {
+        ((k + self.end_padding_record) * self.r).divceil(self.alignment) * self.alignment
     }
 
     #[allow(clippy::too_many_arguments)]
diff --git a/test-rt/suite-unit/src/conv_q.rs b/test-rt/suite-unit/src/conv_q.rs
index 21c0c613a3..7f650a6518 100644
--- a/test-rt/suite-unit/src/conv_q.rs
+++ b/test-rt/suite-unit/src/conv_q.rs
@@ -1201,5 +1201,20 @@ pub fn suite() -> TractResult<TestSuite> {
             raw_output_dt: DatumType::I32,
         },
     );
+    let qp = qp_noop_i8();
+    suite.add(
+        "align",
+        QConvProblem {
+            shape_in: CHW.from_n_c_hw(1, 2, [2]).unwrap(),
+            co: 2,
+            kernel_format: HWIO,
+            group: 2,
+            kernel: tensor3(&[[[0i8], [0]]]),
+            bias: None,
+            data: tensor2(&[[0i8, 0], [0, 0]]),
+            qp,
+            raw_output_dt: DatumType::I8,
+        },
+    );
     Ok(suite)
 }