Mmm def updates

cli/src/main.rs

@@ -587,7 +587,7 @@ fn handle(matches: clap::ArgMatches, probe: Option<&Probe>) -> TractResult<()> {
             for m in tract_linalg::ops().mmm_impls() {
                 println!("{}", Green.paint(format!(" * {}", m.name())));
                 for packings in m.packings() {
-                    println!("   - {} • {}", packings.0, packings.1);
+                    println!("   - {:?} • {:?}", packings.0, packings.1);
                 }
             }
             println!("{}", White.bold().paint("# MatMul kits"));

core/src/ops/einsum/kernel_selection.rs

@@ -110,7 +110,7 @@ pub fn wire_linear(
             kit.weight == weight && kit.accumulator == accumulator && kit.activation == activation
         })
         .min_by_key(|kit| kit.generic_fallback as usize)
-        .with_context(|| format!("No kit found for matmul {a:?} • {b_fact:?}"))?;
+        .with_context(|| format!("No kit found for matmul {weight:?} {accumulator:?} {activation:?}"))?;
     let configs = [kit.item_for_mv(), kit.item_for_squarish()];
     let packed: Box<dyn MMMInputValue> = if let Some(a_payload) = a_as_bqv {
         let packed = kit

data/src/datum.rs

@@ -492,6 +492,7 @@ pub trait Datum:
 {
     fn name() -> &'static str;
     fn datum_type() -> DatumType;
+    fn is<D: Datum>() -> bool;
 }
 
 macro_rules! datum {
@@ -510,6 +511,10 @@ macro_rules! datum {
             fn datum_type() -> DatumType {
                 DatumType::$v
             }
+
+            fn is<D: Datum>() -> bool {
+                Self::datum_type() == D::datum_type()
+            }
         }
     };
 }

linalg/src/arm64/arm64fp16.rs

@@ -50,3 +50,13 @@ pub fn plug(ops: &mut Ops) {
 
 tanh_impl!(f16, arm64fp16_tanh_f16_8n, 8, 8, crate::arm64::has_fp16());
 sigmoid_impl!(f16, arm64fp16_sigmoid_f16_8n, 8, 8, crate::arm64::has_fp16());
+
+#[cfg(test)]
+mod test {
+
+    #[test]
+    fn kits() {
+        let mut ops = crate::generic();
+        super::plug(&mut ops);
+    }
+}

linalg/src/arm64/arm64fp16/panel_extract.rs

@@ -1,16 +1,16 @@
+use super::FP16;
 use crate::frame::block_quant::{PackedBlockQuantFormat, Q4_0};
-use crate::frame::PackedFormat;
+use crate::mmm::Packing;
 use crate::Ops;
 use tract_data::internal::*;
-use super::FP16;
 
 pub fn plug(ops: &mut Ops) {
     ops.panel_extractors.push(packed_64_q40_to_f16.clone());
 }
 
 panel_extractor!(kernel_packed_64_q40_to_f16 as packed_64_q40_to_f16(
     Box::new(PackedBlockQuantFormat::new(&Q4_0, 64, 16, true)),
-    PackedFormat::new(f16::datum_type(), 64, 16)
+    f16::packing(64).align(16)
 ) where(FP16));
 
 #[target_feature(enable = "fp16")]

linalg/src/frame/mmm/kernel.rs

@@ -34,7 +34,6 @@ type Kernel<Acc> = unsafe fn(&[FusedKerSpec<Acc>]) -> isize;
 pub struct DynKernel<const MR: usize, const NR: usize, Acc: LADatum> {
     pub name: String,
     pub kernel: Kernel<Acc>,
-    pub default_packing_alignments: (usize, usize),
     pub packings: Vec<(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)>,
     pub stores: Vec<DatumType>,
     pub supported_predicate: fn() -> bool,
@@ -45,20 +44,18 @@ impl<const MR: usize, const NR: usize, Acc: LADatum> DynKernel<MR, NR, Acc> {
     pub fn new(
         name: &str,
         kernel: Kernel<Acc>,
-        default_packing_alignments: (usize, usize),
+        packing_a: PackedFormat,
+        packing_b: PackedFormat,
     ) -> Self {
         let kernel = DynKernel {
             name: name.to_string(),
             kernel,
             packings: vec![],
             stores: vec![Acc::datum_type()],
             supported_predicate: || true,
-            default_packing_alignments,
             can_fuse: |_| true,
         };
-        let a = kernel.regular_pack_a();
-        let b = kernel.regular_pack_b();
-        kernel.with_packing(a, b)
+        kernel.with_packing(packing_a, packing_b)
     }
 
     pub fn with_platform_condition(mut self, f: fn() -> bool) -> Self {
@@ -77,11 +74,11 @@ impl<const MR: usize, const NR: usize, Acc: LADatum> DynKernel<MR, NR, Acc> {
     }
 
     pub fn regular_pack_a(&self) -> PackedFormat {
-        PackedFormat::new(Acc::datum_type(), MR, self.default_packing_alignments.0)
+        *self.packings[0].0.clone().downcast::<PackedFormat>().unwrap()
     }
 
     pub fn regular_pack_b(&self) -> PackedFormat {
-        PackedFormat::new(Acc::datum_type(), NR, self.default_packing_alignments.1)
+        *self.packings[0].1.clone().downcast::<PackedFormat>().unwrap()
     }
 
     pub fn with_can_fuse(self, can_fuse: fn(&FusedSpec) -> bool) -> Self {

linalg/src/frame/mmm/kit.rs

@@ -187,7 +187,7 @@ impl MMMKit {
                 .1
                 .downcast_ref::<PackedFormat>()
                 .is_some_and(|pf| KitDatumType::from(pf.dt) == self.activation),
-            "Activation packecd mismatch {self:?} {:?}",
+            "Activation packed dt mismatch {self:?} {:?}",
             mmm.packings()[packing].1
         );
         self.items.push(MMMKitItem { mmm, packing, weight_panel_extractor });

linalg/src/frame/mmm/macros.rs

@@ -1,6 +1,7 @@
 macro_rules! MMMExternKernel {
     (
-            $func:ident<$ti:ident>($mr: expr, $nr: expr)@($align_a:expr, $align_b:expr)
+            $func:ident<$ti:ident>($mr: expr, $nr: expr)
+            $(@($align_a:expr, $align_b:expr))?
             $(where($where:expr))?
             $(can_fuse($can_fuse:expr))?
             $(packing[$pnum:literal] = $pid:ident => $packing:expr;)*
@@ -21,7 +22,8 @@ macro_rules! MMMExternKernel {
                 }
             }
 
-            MMMKernel!([<sys_$func>]::rusty as $func<$ti>($mr, $nr)@($align_a, $align_b)
+            MMMKernel!([<sys_$func>]::rusty as $func<$ti>($mr, $nr)
+                $(@($align_a, $align_b))?
                 $(where($where))?
                 $(can_fuse($can_fuse))?
                 $(packing[$pnum] = $pid => $packing;)*
@@ -32,7 +34,8 @@ macro_rules! MMMExternKernel {
 }
 macro_rules! MMMRustKernel {
     (       $func: path =>
-            $id:ident<$ti:ident>($mr: expr, $nr: expr)@($align_a:expr, $align_b:expr)
+            $id:ident<$ti:ident>($mr: expr, $nr: expr)
+            $(@($align_a:expr, $align_b:expr))?
             $(where($where:expr))?
             $(can_fuse($can_fuse:expr))?
             $(packing[$pnum:literal] = $pid:ident => $packing:expr;)*
@@ -49,7 +52,8 @@ macro_rules! MMMRustKernel {
                     $func(op.as_ptr())
                 }
             }
-            MMMKernel!([<sys_$id>]::rusty as $id<$ti>($mr, $nr)@($align_a, $align_b)
+            MMMKernel!([<sys_$id>]::rusty as $id<$ti>($mr, $nr)
+                $(@($align_a, $align_b))?
                 $(where($where))?
                 $(can_fuse($can_fuse))?
                 $(packing[$pnum] = $pid => $packing;)*
@@ -62,7 +66,8 @@ macro_rules! MMMRustKernel {
 macro_rules! MMMKernel {
     (
             $func: path as
-            $id:ident<$ti:ident>($mr: expr, $nr: expr)@($align_a:expr, $align_b:expr)
+            $id:ident<$ti:ident>($mr: expr, $nr: expr)
+            $(@($align_a:expr, $align_b:expr))?
             $(where($where:expr))?
             $(can_fuse($can_fuse:expr))?
             $(packing[$pnum:literal] = $pid:ident => $packing:expr;)*
@@ -75,8 +80,15 @@ macro_rules! MMMKernel {
                     use $crate::mmm::DynKernel;
                     #[allow(unused_imports)]
                     use tract_data::prelude::*;
+                    use $crate::frame::mmm::Packing;
                     #[allow(unused_mut)]
-                    let mut k = DynKernel::<$mr, $nr, $ti>::new(stringify!($id), $func, ($align_a, $align_b));
+                    let (mut packing_a, mut packing_b) = ($ti::packing($mr), $ti::packing($nr));
+                    $(
+                        packing_a = packing_a.align($align_a);
+                        packing_b = packing_b.align($align_b);
+                    )?
+                    #[allow(unused_mut)]
+                    let mut k = DynKernel::<$mr, $nr, $ti>::new(stringify!($id), $func, packing_a, packing_b);
                     $(k = k.with_platform_condition($where);)?
                     $(
                         assert!(k.packings.len() == $pnum);
@@ -102,4 +114,3 @@ macro_rules! MMMKernel {
         }
     };
 }
-

linalg/src/frame/mmm/mod.rs

@@ -30,6 +30,8 @@ pub use kit::*;
 pub use scratch::*;
 pub use storage::*;
 
+pub use pack::Packing;
+
 pub fn no_prefetch(_ptr: *const u8, _len: usize) {}
 
 pub trait MatMatMul: Debug + dyn_clone::DynClone + Send + Sync + std::any::Any {

linalg/src/frame/mmm/pack.rs

@@ -12,7 +12,7 @@ use super::MMMInputFormat;
 pub struct PackedFormat {
     pub dt: DatumType,
     pub r: usize,
-    pub alignment: usize,
+    pub alignment_bytes: usize,
     pub end_padding_record: usize,
 }
 
@@ -47,22 +47,31 @@ impl Display for PackedFormat {
 
 impl Debug for PackedFormat {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        <Self as Display>::fmt(self, f)
+        write!(
+            f,
+            "Packed{:?}[{}]@{}+{}",
+            self.dt, self.r, self.alignment_bytes, self.end_padding_record
+        )
     }
 }
 
 impl PackedFormat {
-    pub const fn new(dt: DatumType, nr: usize, alignment: usize) -> PackedFormat {
-        PackedFormat { dt, r: nr, alignment, end_padding_record: 1 }
+    pub const fn new(dt: DatumType, nr: usize, alignment_bytes: usize) -> PackedFormat {
+        PackedFormat { dt, r: nr, alignment_bytes, end_padding_record: 1 }
     }
 
     pub const fn with_end_padding_record(self, end_padding_record: usize) -> Self {
         PackedFormat { end_padding_record, ..self }
     }
 
+    #[inline]
+    pub fn align(self, alignment: usize) -> Self {
+        Self { alignment_bytes: alignment, ..self }
+    }
+
     #[inline]
     pub fn alignment(&self) -> usize {
-        self.alignment
+        self.alignment_bytes
     }
 
     #[inline]
@@ -101,8 +110,10 @@ impl PackedFormat {
         let panel_bytes = panel_len * t.datum_type().size_of();
         let strides = t.strides();
         unsafe {
-            let mut packed =
-                Blob::new_for_size_and_align(t.datum_type().size_of() * packed_len, self.alignment);
+            let mut packed = Blob::new_for_size_and_align(
+                t.datum_type().size_of() * packed_len,
+                self.alignment_bytes,
+            );
             if cfg!(debug_assertions) {
                 packed.as_bytes_mut().fill(0u8);
             }
@@ -140,8 +151,10 @@ impl PackedFormat {
         let panel_bytes = panel_len * t.datum_type().size_of();
         let strides = t.strides();
         unsafe {
-            let mut packed =
-                Blob::new_for_size_and_align(t.datum_type().size_of() * packed_len, self.alignment);
+            let mut packed = Blob::new_for_size_and_align(
+                t.datum_type().size_of() * packed_len,
+                self.alignment_bytes,
+            );
             if cfg!(debug_assertions) {
                 packed.as_bytes_mut().fill(0u8);
             }
@@ -509,6 +522,16 @@ unsafe fn pack_mn_major<Chunk: Copy>(
     }
 }
 
+pub trait Packing {
+    fn packing(r: usize) -> PackedFormat;
+}
+
+impl<D: Datum> Packing for D {
+    fn packing(r: usize) -> PackedFormat {
+        PackedFormat::new(Self::datum_type(), r, Self::datum_type().alignment())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::ops::Range;