stronger alignment for packing

sonos · Jan 30, 2024 · 7e81938 · 7e81938
1 parent a66ade8
commit 7e81938
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git 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
@@ -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)
 }