tract_linalg/frame/mmm/
input_store.rs

1use downcast_rs::{Downcast, impl_downcast};
2use dyn_clone::DynClone;
3use dyn_eq::DynEq;
4use dyn_hash::DynHash;
5use std::alloc::Layout;
6use std::fmt::{Debug, Display};
7use std::hash::Hash;
8use std::sync::Arc;
9use tract_data::internal::*;
10
11use crate::WeightType;
12
13pub trait MMMInputFormat:
14    Downcast + Debug + DynHash + dyn_eq::DynEq + DynClone + Send + Sync + Display
15{
16    fn prepare_tensor(&self, t: &Tensor, k_axis: usize, mn_axis: usize) -> TractResult<Tensor>;
17    fn prepare_one(
18        &self,
19        t: &Tensor,
20        k_axis: usize,
21        mn_axis: usize,
22    ) -> TractResult<Box<dyn MMMInputValue>>;
23    fn precursor(&self) -> WeightType;
24    /// Round `tensor` through the numeric precision this format stores its data
25    /// in, so a reference matmul can reproduce the kernel's pack-time precision
26    /// loss. Default: identity, for formats that store inputs losslessly.
27    fn simulate_precision_loss(&self, tensor: Tensor) -> TractResult<Tensor> {
28        Ok(tensor)
29    }
30    fn r(&self) -> usize;
31    fn k_alignment(&self) -> usize;
32    fn merge_with<'o, 'a: 'o, 'b: 'o>(
33        &'a self,
34        other: &'b dyn MMMInputFormat,
35    ) -> Option<&'o dyn MMMInputFormat> {
36        if self.dyn_eq(other) { Some(other) } else { None }
37    }
38    fn mem_size(&self, k: TDim, mn: TDim) -> TDim;
39    fn extract_at_mn_f16(
40        &self,
41        data: &EagerPackedInput,
42        mn: usize,
43        slice: &mut [f16],
44    ) -> TractResult<()>;
45    fn extract_at_mn_f32(
46        &self,
47        data: &EagerPackedInput,
48        mn: usize,
49        slice: &mut [f32],
50    ) -> TractResult<()>;
51}
52
53dyn_clone::clone_trait_object!(MMMInputFormat);
54impl_downcast!(MMMInputFormat);
55dyn_hash::hash_trait_object!(MMMInputFormat);
56dyn_eq::eq_trait_object!(MMMInputFormat);
57
58pub trait MMMInputValue:
59    DynClone + Debug + DynHash + dyn_eq::DynEq + Send + Sync + Display + Downcast
60{
61    fn format(&self) -> &dyn MMMInputFormat;
62    fn scratch_panel_buffer_layout(&self) -> Option<Layout>;
63    fn panel_bytes(&self, i: usize, buffer: Option<*mut u8>) -> TractResult<*const u8>;
64    fn panels_count(&self) -> usize {
65        self.mn().divceil(self.format().r())
66    }
67    fn mn(&self) -> usize;
68    fn k(&self) -> usize;
69    fn exotic_fact(&self) -> &dyn ExoticFact;
70
71    fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()>;
72    fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()>;
73}
74dyn_clone::clone_trait_object!(MMMInputValue);
75impl_downcast!(MMMInputValue);
76dyn_hash::hash_trait_object!(MMMInputValue);
77dyn_eq::eq_trait_object!(MMMInputValue);
78
79#[allow(clippy::derived_hash_with_manual_eq)]
80#[derive(Clone, Hash, Debug)]
81pub struct PackedExoticFact {
82    pub format: Box<dyn MMMInputFormat>,
83    pub mn: TDim,
84    pub k: usize,
85}
86
87impl Display for PackedExoticFact {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        write!(f, "Eager {} tensor (mn={} k={})", self.format, self.mn, self.k)
90    }
91}
92
93impl ExoticFact for PackedExoticFact {
94    fn buffer_sizes(&self) -> TVec<TDim> {
95        tvec!(self.format.mem_size(self.k.to_dim(), self.mn.clone()))
96    }
97}
98
99impl PartialEq for PackedExoticFact {
100    fn eq(&self, other: &Self) -> bool {
101        self.format == other.format && self.mn == other.mn && self.k == other.k
102    }
103}
104impl Eq for PackedExoticFact {}
105
106#[derive(Clone, Hash, PartialEq, Eq)]
107pub struct EagerPackedInput {
108    pub fact: PackedExoticFact,
109    pub packed: Arc<Blob>,
110    pub panel_bytes: usize,
111    pub mn: usize,
112}
113
114impl MMMInputValue for EagerPackedInput {
115    fn scratch_panel_buffer_layout(&self) -> Option<Layout> {
116        None
117    }
118    fn panel_bytes(&self, i: usize, _buffer: Option<*mut u8>) -> TractResult<*const u8> {
119        unsafe { Ok(self.packed.as_ptr().add(i * self.panel_bytes)) }
120    }
121    fn k(&self) -> usize {
122        self.fact.k
123    }
124    fn mn(&self) -> usize {
125        self.mn
126    }
127    fn format(&self) -> &dyn MMMInputFormat {
128        &*self.fact.format
129    }
130    fn exotic_fact(&self) -> &dyn ExoticFact {
131        &self.fact
132    }
133    fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()> {
134        ensure!(slice.len() == self.k());
135        ensure!(mn < self.mn());
136        self.fact.format.extract_at_mn_f16(self, mn, slice)
137    }
138    fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()> {
139        ensure!(slice.len() == self.k());
140        ensure!(mn < self.mn());
141        self.fact.format.extract_at_mn_f32(self, mn, slice)
142    }
143}
144
145impl Display for EagerPackedInput {
146    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        (&self.fact as &dyn Display).fmt(f)
148    }
149}
150
151impl Debug for EagerPackedInput {
152    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153        <Self as Display>::fmt(self, f)
154    }
155}
tract_linalg/frame/mmm/input_store.rs

tract_linalg/frame/mmm/
input_store.rs
