add opaque fact in packed tensors

sonos · Oct 14, 2024 · 42fb316 · 42fb316
1 parent 490b636
commit 42fb316
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 17 deletions.
diff --git a/cli/src/run.rs b/cli/src/run.rs
@@ -8,7 +8,7 @@ use fs_err as fs;
 use ndarray_npy::NpzWriter;
 use nu_ansi_term::Color::*;
 use tract_core::ops::cnn::conv::Im2Col;
-use tract_core::ops::matmul::pack::MatMatMulPack;
+use tract_core::ops::matmul::pack::OptMatMulPack;
 use tract_core::tract_data::itertools::izip;
 use tract_hir::internal::*;
 use tract_libcli::tensor::RunParams;
@@ -213,7 +213,7 @@ fn run_regular(
                     }
                     if assert_sane_floats {
                         for (ix, o) in clarified_r.iter().enumerate() {
-                            if node.op_is::<Im2Col>() || node.op_is::<MatMatMulPack>() {
+                            if node.op_is::<Im2Col>() || node.op_is::<OptMatMulPack>() {
                                 continue;
                             }
                             if let Ok(floats) = o.as_slice::<f32>() {

diff --git a/core/src/ops/cnn/conv/conv.rs b/core/src/ops/cnn/conv/conv.rs
@@ -17,7 +17,7 @@ use crate::ops::math::{add, div, mul, sub};
 use crate::ops::math::{Add, Div, Mul, Sub};
 use crate::ops::matmul::optimized::AddMatMulGeometry;
 use crate::ops::matmul::optimized::MapOutputAxisToInput;
-use crate::ops::matmul::pack::MatMatMulPack;
+use crate::ops::matmul::pack::OptMatMulPack;
 use crate::ops::matmul::quant::wire_ensure_q8_flavour;
 use crate::ops::nn::Reduce;
 
@@ -79,7 +79,7 @@ impl Conv {
     ) -> TractResult<OutletId> {
         Ok(model.wire_node(
             format!("{name}.prep_kernel.pack"),
-            MatMatMulPack { packers: vec![format], k_axis: 2, mn_axis: 1 },
+            OptMatMulPack { packers: vec![format], k_axis: 2, mn_axis: 1 },
             &[kernel],
         )?[0])
     }

diff --git a/core/src/ops/einsum/kernel_selection.rs b/core/src/ops/einsum/kernel_selection.rs
@@ -6,7 +6,7 @@ use tract_linalg::mmm::{MMMInputValue, MatMatMul};
 
 use crate::internal::*;
 use crate::ops::matmul::de_block_quant::{BlockQuantFact, BlockQuantValue};
-use crate::ops::matmul::pack::MatMatMulPack;
+use crate::ops::matmul::pack::OptMatMulPack;
 
 use super::optimize::EinSumAnnotatedAsMatMul;
 
@@ -47,13 +47,13 @@ pub fn wire_packing(
         .with_context(|| format!("No packing for {mmm:?} with inputs {a_dt:?} and {b_dt:?}"))?;
     let pa = patch.wire_node(
         format!("{prefix}.pack_a"),
-        MatMatMulPack { k_axis: op.a_k(), mn_axis: op.a_m(), packers: vec![pa.clone()] },
+        OptMatMulPack { k_axis: op.a_k(), mn_axis: op.a_m(), packers: vec![pa.clone()] },
         &[operands[0]],
     )?[0];
 
     let pb = patch.wire_node(
         format!("{prefix}.pack_b"),
-        MatMatMulPack { k_axis: op.b_k(), mn_axis: op.b_n(), packers: vec![pb.clone()] },
+        OptMatMulPack { k_axis: op.b_k(), mn_axis: op.b_n(), packers: vec![pb.clone()] },
         &[operands[1]],
     )?[0];
 
@@ -99,7 +99,7 @@ fn with_block_quant(
             .clone();
         patch
             .node_mut(pb.node)
-            .op_as_mut::<MatMatMulPack>()
+            .op_as_mut::<OptMatMulPack>()
             .context("Expected MatMatMulPack on B")?
             .packers
             .push(alternative_b_packing);
@@ -149,7 +149,7 @@ fn with_block_quant_matmat(
 
     let pb = patch.wire_node(
         format!("{prefix}.pack_b"),
-        MatMatMulPack { k_axis: op.b_k(), mn_axis: op.b_n(), packers: vec![pb.clone()] },
+        OptMatMulPack { k_axis: op.b_k(), mn_axis: op.b_n(), packers: vec![pb.clone()] },
         &[operands[1]],
     )?[0];
 
@@ -198,7 +198,7 @@ fn with_block_quant_matvec(
 
     let pb = patch.wire_node(
         format!("{prefix}.pack_b"),
-        MatMatMulPack { k_axis: op.b_k(), mn_axis: op.b_n(), packers: vec![pb.clone()] },
+        OptMatMulPack { k_axis: op.b_k(), mn_axis: op.b_n(), packers: vec![pb.clone()] },
         &[operands[1]],
     )?[0];
 

diff --git a/core/src/ops/matmul/pack.rs b/core/src/ops/matmul/pack.rs
@@ -5,15 +5,15 @@ use tract_data::TooEarly;
 use tract_linalg::frame::PackedFormat;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct MatMatMulPack {
+pub struct OptMatMulPack {
     pub(crate) packers: Vec<PackedFormat>,
     pub(crate) k_axis: usize,
     pub(crate) mn_axis: usize,
 }
 
-impl Op for MatMatMulPack {
+impl Op for OptMatMulPack {
     fn name(&self) -> Cow<str> {
-        "MatMatMulPack".into()
+        "OptMatMulPack".into()
     }
 
     fn info(&self) -> TractResult<Vec<String>> {
@@ -24,7 +24,7 @@ impl Op for MatMatMulPack {
     impl_op_same_as!();
 }
 
-impl EvalOp for MatMatMulPack {
+impl EvalOp for OptMatMulPack {
     fn is_stateless(&self) -> bool {
         true
     }
@@ -38,9 +38,14 @@ impl EvalOp for MatMatMulPack {
     }
 }
 
-impl TypedOp for MatMatMulPack {
+impl TypedOp for OptMatMulPack {
     fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
-        Ok(tvec!(Opaque::datum_type().fact(self.output_shape(&inputs[0].shape))))
+        let k = inputs[0].shape[self.k_axis].clone();
+        let mn = inputs[0].shape[self.mn_axis].clone();
+        let opaque_fact = PackedOpaqueFact { k, mn, packers: self.packers.clone() };
+        Ok(tvec!(Opaque::datum_type()
+            .fact(self.output_shape(&inputs[0].shape))
+            .with_opaque_fact(opaque_fact)))
     }
 
     fn axes_mapping(
@@ -63,7 +68,7 @@ impl TypedOp for MatMatMulPack {
     as_op!();
 }
 
-impl MatMatMulPack {
+impl OptMatMulPack {
     fn do_eval(&self, session: &SessionState, input: TValue) -> TractResult<TVec<TValue>> {
         unsafe {
             let packer = if self.packers.len() == 1 {
@@ -118,3 +123,16 @@ impl MatMatMulPack {
         packed_shape
     }
 }
+
+#[derive(Hash, Clone, Debug, PartialEq, Eq)]
+pub struct PackedOpaqueFact {
+    pub k: TDim,
+    pub mn: TDim,
+    pub packers: Vec<PackedFormat>,
+}
+
+impl OpaqueFact for PackedOpaqueFact {
+    fn mem_size(&self) -> TDim {
+        self.k.clone() * &self.mn * self.packers[0].dt.size_of()
+    }
+}