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    /// Pack one matmul out of a possibly strided view — one batch of a batched input, which a
18    /// `&Tensor` cannot name. A format whose input is not a strided array of numbers (block
19    /// quant, lazy im2col) says so here rather than being silently unreachable.
20    fn prepare_one_view(
21        &self,
22        t: &TensorView,
23        k_axis: usize,
24        mn_axis: usize,
25    ) -> TractResult<Box<dyn MMMInputValue>>;
26    fn prepare_one(
27        &self,
28        t: &Tensor,
29        k_axis: usize,
30        mn_axis: usize,
31    ) -> TractResult<Box<dyn MMMInputValue>> {
32        self.prepare_one_view(&t.view(), k_axis, mn_axis)
33    }
34    fn precursor(&self) -> WeightType;
35    /// Round `tensor` through the numeric precision this format stores its data
36    /// in, so a reference matmul can reproduce the kernel's pack-time precision
37    /// loss. Default: identity, for formats that store inputs losslessly.
38    fn simulate_precision_loss(&self, tensor: Tensor) -> TractResult<Tensor> {
39        Ok(tensor)
40    }
41    fn r(&self) -> usize;
42    fn k_alignment(&self) -> usize;
43    fn merge_with<'o, 'a: 'o, 'b: 'o>(
44        &'a self,
45        other: &'b dyn MMMInputFormat,
46    ) -> Option<&'o dyn MMMInputFormat> {
47        if self.dyn_eq(other) { Some(other) } else { None }
48    }
49    fn mem_size(&self, k: TDim, mn: TDim) -> TDim;
50    fn extract_at_mn_f16(
51        &self,
52        data: &EagerPackedInput,
53        mn: usize,
54        slice: &mut [f16],
55    ) -> TractResult<()>;
56    fn extract_at_mn_f32(
57        &self,
58        data: &EagerPackedInput,
59        mn: usize,
60        slice: &mut [f32],
61    ) -> TractResult<()>;
62}
63
64dyn_clone::clone_trait_object!(MMMInputFormat);
65impl_downcast!(MMMInputFormat);
66dyn_hash::hash_trait_object!(MMMInputFormat);
67dyn_eq::eq_trait_object!(MMMInputFormat);
68
69pub trait MMMInputValue:
70    DynClone + Debug + DynHash + dyn_eq::DynEq + Send + Sync + Display + Downcast
71{
72    fn format(&self) -> &dyn MMMInputFormat;
73    fn scratch_panel_buffer_layout(&self) -> Option<Layout>;
74    fn panel_bytes(&self, i: usize, buffer: Option<*mut u8>) -> TractResult<*const u8>;
75    fn panels_count(&self) -> usize {
76        self.mn().divceil(self.format().r())
77    }
78    fn mn(&self) -> usize;
79    fn k(&self) -> usize;
80    fn exotic_fact(&self) -> &dyn ExoticFact;
81
82    fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()>;
83    fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()>;
84}
85dyn_clone::clone_trait_object!(MMMInputValue);
86impl_downcast!(MMMInputValue);
87dyn_hash::hash_trait_object!(MMMInputValue);
88dyn_eq::eq_trait_object!(MMMInputValue);
89
90#[allow(clippy::derived_hash_with_manual_eq)]
91#[derive(Clone, Hash, Debug)]
92pub struct PackedExoticFact {
93    pub format: Box<dyn MMMInputFormat>,
94    pub mn: TDim,
95    pub k: usize,
96}
97
98impl Display for PackedExoticFact {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        write!(f, "Eager {} tensor (mn={} k={})", self.format, self.mn, self.k)
101    }
102}
103
104impl ExoticFact for PackedExoticFact {
105    fn buffer_sizes(&self) -> TVec<TDim> {
106        tvec!(self.format.mem_size(self.k.to_dim(), self.mn.clone()))
107    }
108}
109
110impl PartialEq for PackedExoticFact {
111    fn eq(&self, other: &Self) -> bool {
112        self.format == other.format && self.mn == other.mn && self.k == other.k
113    }
114}
115impl Eq for PackedExoticFact {}
116
117#[derive(Clone, Hash, PartialEq, Eq)]
118pub struct EagerPackedInput {
119    pub fact: PackedExoticFact,
120    pub packed: Arc<Blob>,
121    pub panel_bytes: usize,
122    pub mn: usize,
123}
124
125impl MMMInputValue for EagerPackedInput {
126    fn scratch_panel_buffer_layout(&self) -> Option<Layout> {
127        None
128    }
129    fn panel_bytes(&self, i: usize, _buffer: Option<*mut u8>) -> TractResult<*const u8> {
130        unsafe { Ok(self.packed.as_ptr().add(i * self.panel_bytes)) }
131    }
132    fn k(&self) -> usize {
133        self.fact.k
134    }
135    fn mn(&self) -> usize {
136        self.mn
137    }
138    fn format(&self) -> &dyn MMMInputFormat {
139        &*self.fact.format
140    }
141    fn exotic_fact(&self) -> &dyn ExoticFact {
142        &self.fact
143    }
144    fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()> {
145        ensure!(slice.len() == self.k());
146        ensure!(mn < self.mn());
147        self.fact.format.extract_at_mn_f16(self, mn, slice)
148    }
149    fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()> {
150        ensure!(slice.len() == self.k());
151        ensure!(mn < self.mn());
152        self.fact.format.extract_at_mn_f32(self, mn, slice)
153    }
154}
155
156impl Display for EagerPackedInput {
157    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158        (&self.fact as &dyn Display).fmt(f)
159    }
160}
161
162impl Debug for EagerPackedInput {
163    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164        <Self as Display>::fmt(self, f)
165    }
166}
167
168#[cfg(test)]
169pub mod view_check {
170    use super::*;
171
172    /// Packing one batch of a batched tensor through a strided view must give the same bytes as
173    /// packing that batch alone: what a dynamic packing op relies on, and what an unimplemented
174    /// [`MMMInputFormat::prepare_one_view`] used to break.
175    pub fn view_matches_tensor(
176        format: &dyn MMMInputFormat,
177        k: usize,
178        mn: usize,
179    ) -> TractResult<()> {
180        let WeightType::Plain(dt) = format.precursor() else {
181            bail!("{format} does not pack a plain tensor")
182        };
183        let mut parent = Tensor::zero_dt(dt, &[2, k, mn])?;
184        parent.as_bytes_mut().iter_mut().enumerate().for_each(|(ix, b)| *b = (ix % 101) as u8 + 1);
185        let batch = parent.slice(0, 1, 2)?.into_shape(&[k, mn])?;
186        let from_tensor = format.prepare_one(&batch, 0, 1)?;
187        let offset = parent.strides()[0] * dt.size_of() as isize;
188        let view =
189            unsafe { TensorView::from_bytes(&parent, offset, parent.shape(), parent.strides()) };
190        let from_view = format.prepare_one_view(&view, 1, 2)?;
191        let bytes = |v: &dyn MMMInputValue| -> TractResult<Arc<Blob>> {
192            Ok(v.downcast_ref::<EagerPackedInput>()
193                .with_context(|| format!("{format} did not pack to an EagerPackedInput"))?
194                .packed
195                .clone())
196        };
197        ensure!(
198            *bytes(&*from_tensor)? == *bytes(&*from_view)?,
199            "{format}: packing k={k} mn={mn} from a view differs from packing it from a tensor"
200        );
201        Ok(())
202    }
203}
tract_linalg/frame/mmm/input_store.rs

tract_linalg/frame/mmm/
input_store.rs
