tract_linalg/frame/mmm/
mod.rs

1#[macro_use]
2mod macros;
3
4pub mod cost_model;
5#[macro_use]
6pub(crate) mod fuse;
7pub(crate) mod input_store;
8pub(crate) mod kernel;
9#[macro_use]
10pub(crate) mod panel_extract;
11mod scratch;
12mod storage;
13
14#[cfg(test)]
15#[macro_use]
16pub mod tests;
17
18use crate::multithread::Executor;
19#[cfg(feature = "multithread-mm")]
20use rayon::prelude::*;
21use std::borrow::Cow;
22use std::cmp::Ordering;
23use std::fmt::Debug;
24use tract_data::internal::*;
25
26pub use cost_model::*;
27pub use fuse::*;
28pub use input_store::*;
29pub use kernel::*;
30pub use panel_extract::*;
31pub use scratch::*;
32pub use storage::*;
33
34pub fn no_prefetch(_ptr: *const u8, _len: usize) {}
35
36#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
37pub enum ImplementationQuality {
38    /// Individual operations are emulated by individual conversion (f16->f32->f16)
39    Dreadful,
40    /// Rust scalar operation (with whatever optimisation the compiler manages)
41    Generic,
42    /// Implicit vectorization (e.g. Rust code, some unrolled loops, explicit template instantiations for small constant)
43    RustOptimized,
44    /// Explicit vectorization (e.g. intrinsics vector code)
45    TargetOptimized,
46    /// Hand optimized (assembly)
47    ManuallyOptimized,
48}
49
50impl ImplementationQuality {
51    pub fn best_to_worst() -> &'static [ImplementationQuality] {
52        use ImplementationQuality::*;
53        &[ManuallyOptimized, TargetOptimized, RustOptimized, Generic, Dreadful]
54    }
55
56    pub fn cost(&self) -> usize {
57        ImplementationQuality::best_to_worst().iter().position(|x| x == self).unwrap()
58    }
59}
60
61impl PartialOrd for ImplementationQuality {
62    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
63        Some(usize::from(*self).cmp(&usize::from(*other)))
64    }
65}
66
67impl From<ImplementationQuality> for usize {
68    fn from(value: ImplementationQuality) -> Self {
69        value.cost()
70    }
71}
72
73pub trait MatMatMul: Debug + dyn_clone::DynClone + Send + Sync + std::any::Any {
74    fn name(&self) -> &str;
75    fn mr(&self) -> usize;
76    fn nr(&self) -> usize;
77
78    fn quality(&self) -> ImplementationQuality;
79    fn dynamic_boost(&self) -> isize;
80
81    #[allow(clippy::type_complexity)]
82    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)];
83
84    fn internal_type(&self) -> DatumType;
85
86    unsafe fn c_view(&self, m_axis: Option<usize>, n_axis: Option<usize>) -> OutputStoreSpec;
87    unsafe fn c_from_data_and_strides(
88        &self,
89        item_size: usize,
90        row_stride: isize,
91        col_stride: isize,
92    ) -> OutputStoreSpec;
93
94    fn can_fuse(&self, spec: &FusedSpec) -> bool;
95
96    fn stores(&self) -> Cow<'_, [DatumType]>;
97
98    unsafe fn run(&self, m: usize, n: usize, non_linear: &[FusedSpec]) -> TractResult<()> {
99        unsafe {
100            let mut scratch = self.allocate_scratch_space();
101            self.run_with_scratch_space(m, n, &mut *scratch, non_linear)
102        }
103    }
104
105    unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace>;
106    unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool;
107    unsafe fn run_with_scratch_space(
108        &self,
109        m: usize,
110        n: usize,
111        scratch: &mut dyn ScratchSpace,
112        non_linear: &[FusedSpec],
113    ) -> TractResult<()>;
114}
115
116dyn_clone::clone_trait_object!(MatMatMul);
117
118impl PartialEq for Box<dyn MatMatMul> {
119    fn eq(&self, other: &Box<dyn MatMatMul>) -> bool {
120        self.name() == other.name()
121    }
122}
123
124impl std::hash::Hash for Box<dyn MatMatMul> {
125    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
126        self.name().hash(state)
127    }
128}
129
130impl<K: MatMatMulKer> MatMatMul for K {
131    fn name(&self) -> &str {
132        self.name()
133    }
134    fn mr(&self) -> usize {
135        self.mr()
136    }
137    fn nr(&self) -> usize {
138        self.nr()
139    }
140
141    fn quality(&self) -> ImplementationQuality {
142        MatMatMulKer::quality(self)
143    }
144
145    fn dynamic_boost(&self) -> isize {
146        MatMatMulKer::dynamic_boost(self)
147    }
148
149    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
150        self.packings()
151    }
152
153    fn internal_type(&self) -> DatumType {
154        K::Acc::datum_type()
155    }
156
157    fn can_fuse(&self, spec: &FusedSpec) -> bool {
158        self.can_fuse(spec)
159    }
160
161    unsafe fn c_view(&self, m_axis: Option<usize>, n_axis: Option<usize>) -> OutputStoreSpec {
162        OutputStoreSpec::View { m_axis, n_axis, mr: self.mr(), nr: self.nr() }
163    }
164
165    unsafe fn c_from_data_and_strides(
166        &self,
167        item_size: usize,
168        row_stride: isize,
169        col_stride: isize,
170    ) -> OutputStoreSpec {
171        OutputStoreSpec::Strides {
172            row_byte_stride: row_stride * item_size as isize,
173            col_byte_stride: col_stride * item_size as isize,
174            mr: self.mr(),
175            nr: self.nr(),
176        }
177    }
178
179    fn stores(&self) -> Cow<'_, [DatumType]> {
180        self.stores()
181    }
182
183    unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace> {
184        Box::<ScratchSpaceImpl<K::Acc>>::default()
185    }
186
187    unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool {
188        scratch.downcast_ref::<ScratchSpaceImpl<K::Acc>>().is_some()
189    }
190
191    unsafe fn run_with_scratch_space(
192        &self,
193        m: usize,
194        n: usize,
195        scratch: &mut dyn ScratchSpace,
196        non_linear: &[FusedSpec],
197    ) -> TractResult<()> {
198        unsafe {
199            let scratch = scratch
200                .downcast_mut::<ScratchSpaceImpl<K::Acc>>()
201                .context("Wrong scratch space type")?;
202            scratch.prepare(self, m, n, non_linear)?;
203            if n == 1 && self.nr() == 1 {
204                run_with_scratch_space_vec(self, m, scratch, non_linear)
205            } else {
206                let (mut prefer_col, mut prefer_row) = (0, 0);
207                for uop in non_linear.iter() {
208                    if let Some(col) = uop.prefer_col_outer() {
209                        prefer_col = col as usize;
210                        prefer_row = (!col) as usize;
211                    }
212                }
213                if prefer_col > prefer_row {
214                    run_with_scratch_space_col_outer(self, m, n, scratch, non_linear)
215                } else {
216                    run_with_scratch_space_row_outer(self, m, n, scratch, non_linear)
217                }
218            }
219        }
220    }
221}
222
223unsafe fn run_with_scratch_space_vec<K: MatMatMulKer>(
224    ker: &K,
225    m: usize,
226    scratch: &mut ScratchSpaceImpl<K::Acc>,
227    non_linear: &[FusedSpec],
228) -> TractResult<()> {
229    unsafe {
230        match crate::multithread::current_tract_executor() {
231            Executor::SingleThread => {
232                for ia in 0..m.divceil(ker.mr()) {
233                    scratch.run(ker, non_linear, ia, 0)?;
234                }
235                Ok(())
236            }
237            #[cfg(feature = "multithread-mm")]
238            Executor::MultiThread(pool) => pool.install(|| {
239                (0..m.div_ceil(ker.mr()))
240                    .into_par_iter()
241                    .try_for_each(|ia| scratch.run(ker, non_linear, ia, 0))
242            }),
243        }
244    }
245}
246
247unsafe fn run_with_scratch_space_col_outer<K: MatMatMulKer>(
248    ker: &K,
249    m: usize,
250    n: usize,
251    scratch: &mut ScratchSpaceImpl<K::Acc>,
252    non_linear: &[FusedSpec],
253) -> TractResult<()> {
254    unsafe {
255        match crate::multithread::current_tract_executor() {
256            Executor::SingleThread => {
257                for ib in 0..n.divceil(ker.nr()) {
258                    for ia in 0..m.divceil(ker.mr()) {
259                        scratch.run(ker, non_linear, ia, ib)?;
260                    }
261                }
262                Ok(())
263            }
264            #[cfg(feature = "multithread-mm")]
265            Executor::MultiThread(pool) => pool.install(|| {
266                (0..n.div_ceil(ker.nr())).into_par_iter().try_for_each(|ib| {
267                    for ia in 0..m.divceil(ker.mr()) {
268                        scratch.run(ker, non_linear, ia, ib)?;
269                    }
270                    Ok(())
271                })
272            }),
273        }
274    }
275}
276
277unsafe fn run_with_scratch_space_row_outer<K: MatMatMulKer>(
278    ker: &K,
279    m: usize,
280    n: usize,
281    scratch: &mut ScratchSpaceImpl<K::Acc>,
282    non_linear: &[FusedSpec],
283) -> TractResult<()> {
284    unsafe {
285        match crate::multithread::current_tract_executor() {
286            Executor::SingleThread => {
287                for ia in 0..m.divceil(ker.mr()) {
288                    for ib in 0..n.divceil(ker.nr()) {
289                        scratch.run(ker, non_linear, ia, ib)?;
290                    }
291                }
292                Ok(())
293            }
294            #[cfg(feature = "multithread-mm")]
295            Executor::MultiThread(pool) => pool.install(|| {
296                pool.install(|| {
297                    (0..m.div_ceil(ker.mr())).into_par_iter().try_for_each(|ia| {
298                        for ib in 0..n.divceil(ker.nr()) {
299                            scratch.run(ker, non_linear, ia, ib)?;
300                        }
301                        Ok(())
302                    })
303                })
304            }),
305        }
306    }
307}
tract_linalg/frame/mmm/mod.rs

tract_linalg/frame/mmm/
mod.rs
