less flexible types to reduce matmul overhead

sonos · Oct 11, 2024 · 7e6a457 · 7e6a457
1 parent 9a1cb28
commit 7e6a457
Show file tree

Hide file tree

Showing 8 changed files with 17 additions and 18 deletions.
diff --git a/core/src/ops/matmul/optimized.rs b/core/src/ops/matmul/optimized.rs
@@ -34,7 +34,7 @@ impl ProtoFusedSpec {
         use ProtoFusedSpec::*;
         match self {
             AddMatMul { geo, packings: packing, .. } => {
-                let (a, b) = mmm.packings()[packing[scenario]];
+                let (a, b) = &mmm.packings()[packing[scenario]];
                 format!("matmul(k={}, {a:?}•{b:?})", geo.k)
             }
             BinScalar(_, op) => format!("scalar{op:?}"),
@@ -72,7 +72,7 @@ impl ProtoFusedSpec {
                 let b = b.as_slice::<Opaque>().unwrap()[0]
                     .downcast_ref::<Box<dyn MMMInputValue>>()
                     .unwrap();
-                let (a_packing, b_packing) = mmm.packings()[packings[scenario]];
+                let (a_packing, b_packing) = &mmm.packings()[packings[scenario]];
                 let pa = if a_packing.same_as(a.format()) {
                     AsInputValue::Borrowed(&**a)
                 } else if a_packing.is::<PackedFormat>()
@@ -155,7 +155,7 @@ impl ProtoFusedSpec {
                 let b = b.as_slice::<Opaque>().unwrap()[0]
                     .downcast_ref::<Box<dyn MMMInputValue>>()
                     .unwrap();
-                let (a_packing, b_packing) = mmm.packings()[packings[scenario]];
+                let (a_packing, b_packing) = &mmm.packings()[packings[scenario]];
                 let pa = if a_packing.same_as(a.format()) {
                     AsInputValue::Borrowed(&**a)
                 } else if a_packing.is::<PackedFormat>()

diff --git a/linalg/benches/mat_vec.rs b/linalg/benches/mat_vec.rs
@@ -16,7 +16,7 @@ fn mat_vec_mul(c: &mut Criterion) {
                 &(m, k),
                 |be, (&m, &k)| {
                     let mmm = tract_linalg::ops().mmm(F32, Some(m), Some(k), Some(1)).unwrap();
-                    let packing = mmm.packings()[0];
+                    let packing = &mmm.packings()[0];
                     let a = Tensor::zero::<f32>(&[m, k]).unwrap();
                     let pa = packing.0.prepare_tensor(&a, 1, 0).unwrap();
                     let b = Tensor::zero::<f32>(&[k, 1]).unwrap();

diff --git a/linalg/benches/mm_for_inception.rs b/linalg/benches/mm_for_inception.rs
@@ -11,7 +11,7 @@ fn mat_mul_smmm(be: &mut criterion::Bencher, &(m, k, n): &(usize, usize, usize))
         let mmm = tract_linalg::ops().mmm(F32, Some(m), Some(k), Some(n)).unwrap();
         let a = Tensor::zero::<f32>(&[m, k]).unwrap();
         let b = Tensor::zero::<f32>(&[k, n]).unwrap();
-        let packing = mmm.packings()[0];
+        let packing = &mmm.packings()[0];
         let pa = packing.0.prepare_tensor(&a, 1, 0).unwrap();
         let pb = packing.1.prepare_tensor(&b, 0, 1).unwrap();
 

diff --git a/linalg/benches/utils.rs b/linalg/benches/utils.rs
@@ -82,7 +82,7 @@ pub fn mat_mat_with_mm(
 ) {
     let a = Tensor::zero_dt(dt, &[m, k]).unwrap();
     let b = Tensor::zero_dt(dt, &[k, n]).unwrap();
-    let packing = mmm.packings()[0];
+    let packing = &mmm.packings()[0];
     let pa = packing.0.prepare_tensor(&a, 1, 0).unwrap();
     let pb = packing.1.prepare_tensor(&b, 0, 1).unwrap();
     unsafe {

diff --git a/linalg/src/frame/mmm/kernel.rs b/linalg/src/frame/mmm/kernel.rs
@@ -1,5 +1,4 @@
 use pack::PackedFormat;
-use tract_itertools::Itertools;
 
 use super::*;
 use std::borrow::Cow;
@@ -14,7 +13,7 @@ pub trait MatMatMulKer: Clone + Debug + Send + Sync + 'static {
     fn mr(&self) -> usize;
     fn nr(&self) -> usize;
 
-    fn packings(&self) -> Cow<[(&dyn MMMInputFormat, &dyn MMMInputFormat)]>;
+    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)];
     fn stores(&self) -> Cow<[DatumType]>;
 
     #[allow(unused_variables)]
@@ -46,7 +45,7 @@ impl<const MR: usize, const NR: usize, Acc: LADatum> DynKernel<MR, NR, Acc> {
         name: &str,
         kernel: Kernel<Acc>,
         default_packing_alignments: (usize, usize),
-    ) -> Self {
+        ) -> Self {
         let kernel = DynKernel {
             name: name.to_string(),
             kernel,
@@ -130,11 +129,11 @@ impl<const MR: usize, const NR: usize, Acc: LADatum> MatMatMulKer for DynKernel<
         unsafe { (self.kernel)(op) }
     }
 
-    fn packings(&self) -> Cow<[(&dyn MMMInputFormat, &dyn MMMInputFormat)]> {
-        Cow::Owned(self.packings.iter().map(|p| (&*p.0, &*p.1)).collect_vec())
+    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
+        &self.packings
     }
 
     fn stores(&self) -> Cow<[DatumType]> {
         Cow::Borrowed(&self.stores)
     }
-}
+    }
diff --git a/linalg/src/frame/mmm/mod.rs b/linalg/src/frame/mmm/mod.rs
@@ -34,7 +34,7 @@ pub trait MatMatMul: Debug + dyn_clone::DynClone + Send + Sync + std::any::Any {
     fn mr(&self) -> usize;
     fn nr(&self) -> usize;
 
-    fn packings(&self) -> Cow<[(&dyn MMMInputFormat, &dyn MMMInputFormat)]>;
+    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)];
 
     fn internal_type(&self) -> DatumType;
 
@@ -89,7 +89,7 @@ impl<K: MatMatMulKer> MatMatMul for K {
         self.nr()
     }
 
-    fn packings(&self) -> Cow<[(&dyn MMMInputFormat, &dyn MMMInputFormat)]> {
+    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
         self.packings()
     }
 

diff --git a/linalg/src/frame/mmm/tests/packed_packed.rs b/linalg/src/frame/mmm/tests/packed_packed.rs
@@ -254,7 +254,7 @@ impl<K: MatMatMulKer> PackedPackedProblem<K> {
     }
 
     pub fn padded_inputs(&self) -> TractResult<(Tensor, Tensor)> {
-        let (pack_a, pack_b) = self.ker.packings()[self.packing];
+        let (pack_a, pack_b) = &self.ker.packings()[self.packing];
         assert!(pack_b.k_alignment() == 1);
         let (m, k, n) = self.mkn();
         let k_aligned = k.next_multiple_of(pack_a.k_alignment());
@@ -283,7 +283,7 @@ impl<K: MatMatMulKer> PackedPackedProblem<K> {
 
     pub fn reference(&self) -> TractResult<Tensor> {
         let (m, k, n) = self.mkn();
-        let pack_a = self.ker.packings()[self.packing].0;
+        let pack_a = &self.ker.packings()[self.packing].0;
         let (mut a, b) = self.padded_inputs()?;
         let k_aligned = k.next_multiple_of(pack_a.k_alignment());
         if let Some(pbqf) = pack_a.downcast_ref::<PackedBlockQuantFormat>() {
@@ -310,7 +310,7 @@ impl<K: MatMatMulKer> PackedPackedProblem<K> {
 
     pub fn run(&self) -> TractResult<Tensor> {
         let (m, k, n) = self.mkn();
-        let (pack_a, pack_b) = self.ker.packings()[self.packing];
+        let (pack_a, pack_b) = &self.ker.packings()[self.packing];
         assert!(pack_b.k_alignment() == 1);
         let k_aligned = k.next_multiple_of(pack_a.k_alignment());
 

diff --git a/linalg/tests/virtual_im2col.rs b/linalg/tests/virtual_im2col.rs
@@ -151,7 +151,7 @@ impl ConvProblem {
         let mmm = tract_linalg::ops().mmm(F32, Some(m), Some(k), Some(n)).unwrap();
         let output = Tensor::zero::<f32>(&internal_output_shape)?;
         let reshaped_filters = self.filters.clone().into_shape(&[k, m])?;
-        let (a_pack, b_pack) = mmm.packings()[0];
+        let (a_pack, b_pack) = &mmm.packings()[0];
         let a = a_pack.prepare_tensor(&reshaped_filters, 0, 1)?;
         unsafe {
             let im2col: Box<dyn MMMInputValue> = if self.lazy_im2col {