AVX512 kit

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/build.rs

@@ -269,6 +269,7 @@ fn preprocess_file(
     }
     .to_owned();
     let long = if msvc { "dd" } else { ".long" };
+    let quad = if msvc { "dq" } else { ".quad" };
     let g = if os == "macos" || os == "ios" { "_" } else { "" };
     // note: use .align with bytes instead of p2align since they both use direct bytes.
     let align = if msvc { "align" } else { ".align" };
@@ -281,6 +282,7 @@ fn preprocess_file(
         "G": g,
         "suffix": suffix,
         "long": long,
+        "quad": quad,
         "jump_table": jump_table(),
         "align": align,
         "offset": if msvc { "offset" } else { "rip + "},

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/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,27 @@ 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]
@@ -102,7 +107,7 @@ impl PackedFormat {
         let strides = t.strides();
         unsafe {
             let mut packed =
-                Blob::new_for_size_and_align(t.datum_type().size_of() * packed_len, self.alignment);
+                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);
             }
@@ -141,7 +146,7 @@ impl PackedFormat {
         let strides = t.strides();
         unsafe {
             let mut packed =
-                Blob::new_for_size_and_align(t.datum_type().size_of() * packed_len, self.alignment);
+                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 +514,21 @@ 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 {
+            dt: Self::datum_type(),
+            r,
+            alignment_bytes: Self::datum_type().alignment(),
+            end_padding_record: 0,
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::ops::Range;

linalg/src/frame/mmm/tests/fuse.rs

@@ -172,7 +172,7 @@ where
     let v = c.to_vec();
     let c = mmm_stride_storage(&v, ker.nr());
     let mut ops = ops.to_vec();
-    ops.insert(0, FusedKerSpec::AddUnicast(c));
+    ops.insert(0, FusedKerSpec::AddUnicast(c)); // FIXME
     ops.insert(0, FusedKerSpec::Clear);
     ops.push(FusedKerSpec::Store(c));
     ops.push(FusedKerSpec::Done);

linalg/src/frame/mmm/tests/packed_packed.rs

@@ -22,7 +22,7 @@ macro_rules! mmm_packed_packed_tests {
             #[allow(unused_imports)]
             use $crate::frame::mmm::tests::packed_packed::*;
 
-            mod fuse {
+            mod packed_packed {
                 use super::*;
 
                 proptest::proptest! {
@@ -56,6 +56,28 @@ macro_rules! mmm_packed_packed_tests {
                     t((1..=$ker.mr() as i64).map(|x| x as f32).collect_vec(), vec![1f32; $ker.nr()])
                 }
 
+                #[test]
+                fn packed_packed_a_scale_times_2_left() -> TractResult<()> {
+                    t(
+                        (1..=2 * $ker.mr() as i64).map(|x| x as f32).collect_vec(),
+                        vec![1f32; $ker.nr()]
+                            .into_iter()
+                            .chain(vec![0f32; $ker.nr()].into_iter())
+                            .collect_vec(),
+                    )
+                }
+
+                #[test]
+                fn packed_packed_a_scale_times_2_right() -> TractResult<()> {
+                    t(
+                        (1..=2 * $ker.mr() as i64).map(|x| x as f32).collect_vec(),
+                        vec![0f32; $ker.nr()]
+                            .into_iter()
+                            .chain(vec![1f32; $ker.nr()].into_iter())
+                            .collect_vec(),
+                    )
+                }
+
                 #[test]
                 fn packed_packed_a_scale_times_2() -> TractResult<()> {
                     t(