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        let mut scratch = self.allocate_scratch_space();
100        self.run_with_scratch_space(m, n, &mut *scratch, non_linear)
101    }
102
103    unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace>;
104    unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool;
105    unsafe fn run_with_scratch_space(
106        &self,
107        m: usize,
108        n: usize,
109        scratch: &mut dyn ScratchSpace,
110        non_linear: &[FusedSpec],
111    ) -> TractResult<()>;
112}
113
114dyn_clone::clone_trait_object!(MatMatMul);
115
116impl PartialEq for Box<dyn MatMatMul> {
117    fn eq(&self, other: &Box<dyn MatMatMul>) -> bool {
118        self.as_ref().type_id() == other.as_ref().type_id()
119    }
120}
121
122impl std::hash::Hash for Box<dyn MatMatMul> {
123    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
124        self.as_ref().type_id().hash(state)
125    }
126}
127
128impl<K: MatMatMulKer> MatMatMul for K {
129    fn name(&self) -> &str {
130        self.name()
131    }
132    fn mr(&self) -> usize {
133        self.mr()
134    }
135    fn nr(&self) -> usize {
136        self.nr()
137    }
138
139    fn quality(&self) -> ImplementationQuality {
140        MatMatMulKer::quality(self)
141    }
142
143    fn dynamic_boost(&self) -> isize {
144        MatMatMulKer::dynamic_boost(self)
145    }
146
147    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
148        self.packings()
149    }
150
151    fn internal_type(&self) -> DatumType {
152        K::Acc::datum_type()
153    }
154
155    fn can_fuse(&self, spec: &FusedSpec) -> bool {
156        self.can_fuse(spec)
157    }
158
159    unsafe fn c_view(&self, m_axis: Option<usize>, n_axis: Option<usize>) -> OutputStoreSpec {
160        OutputStoreSpec::View { m_axis, n_axis, mr: self.mr(), nr: self.nr() }
161    }
162
163    unsafe fn c_from_data_and_strides(
164        &self,
165        item_size: usize,
166        row_stride: isize,
167        col_stride: isize,
168    ) -> OutputStoreSpec {
169        OutputStoreSpec::Strides {
170            row_byte_stride: row_stride * item_size as isize,
171            col_byte_stride: col_stride * item_size as isize,
172            mr: self.mr(),
173            nr: self.nr(),
174        }
175    }
176
177    fn stores(&self) -> Cow<[DatumType]> {
178        self.stores()
179    }
180
181    unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace> {
182        Box::<ScratchSpaceImpl<K::Acc>>::default()
183    }
184
185    unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool {
186        scratch.downcast_ref::<ScratchSpaceImpl<K::Acc>>().is_some()
187    }
188
189    unsafe fn run_with_scratch_space(
190        &self,
191        m: usize,
192        n: usize,
193        scratch: &mut dyn ScratchSpace,
194        non_linear: &[FusedSpec],
195    ) -> TractResult<()> {
196        let scratch = scratch
197            .downcast_mut::<ScratchSpaceImpl<K::Acc>>()
198            .context("Wrong scratch space type")?;
199        scratch.prepare(self, m, n, non_linear)?;
200        if n == 1 && self.nr() == 1 {
201            run_with_scratch_space_vec(self, m, scratch, non_linear)
202        } else {
203            let (mut prefer_col, mut prefer_row) = (0, 0);
204            for uop in non_linear.iter() {
205                if let Some(col) = uop.prefer_col_outer() {
206                    prefer_col = col as usize;
207                    prefer_row = (!col) as usize;
208                }
209            }
210            if prefer_col > prefer_row {
211                run_with_scratch_space_col_outer(self, m, n, scratch, non_linear)
212            } else {
213                run_with_scratch_space_row_outer(self, m, n, scratch, non_linear)
214            }
215        }
216    }
217}
218
219unsafe fn run_with_scratch_space_vec<K: MatMatMulKer>(
220    ker: &K,
221    m: usize,
222    scratch: &mut ScratchSpaceImpl<K::Acc>,
223    non_linear: &[FusedSpec],
224) -> TractResult<()> {
225    match crate::multithread::current_tract_executor() {
226        Executor::SingleThread => {
227            for ia in 0..m.divceil(ker.mr()) {
228                scratch.run(ker, non_linear, ia, 0)?;
229            }
230            Ok(())
231        }
232        #[cfg(feature = "multithread-mm")]
233        Executor::MultiThread(pool) => pool.install(|| {
234            (0..m.div_ceil(ker.mr()))
235                .into_par_iter()
236                .try_for_each(|ia| scratch.run(ker, non_linear, ia, 0))
237        }),
238    }
239}
240
241unsafe fn run_with_scratch_space_col_outer<K: MatMatMulKer>(
242    ker: &K,
243    m: usize,
244    n: usize,
245    scratch: &mut ScratchSpaceImpl<K::Acc>,
246    non_linear: &[FusedSpec],
247) -> TractResult<()> {
248    match crate::multithread::current_tract_executor() {
249        Executor::SingleThread => {
250            for ib in 0..n.divceil(ker.nr()) {
251                for ia in 0..m.divceil(ker.mr()) {
252                    scratch.run(ker, non_linear, ia, ib)?;
253                }
254            }
255            Ok(())
256        }
257        #[cfg(feature = "multithread-mm")]
258        Executor::MultiThread(pool) => pool.install(|| {
259            (0..n.div_ceil(ker.nr())).into_par_iter().try_for_each(|ib| {
260                for ia in 0..m.divceil(ker.mr()) {
261                    scratch.run(ker, non_linear, ia, ib)?;
262                }
263                Ok(())
264            })
265        }),
266    }
267}
268
269unsafe fn run_with_scratch_space_row_outer<K: MatMatMulKer>(
270    ker: &K,
271    m: usize,
272    n: usize,
273    scratch: &mut ScratchSpaceImpl<K::Acc>,
274    non_linear: &[FusedSpec],
275) -> TractResult<()> {
276    match crate::multithread::current_tract_executor() {
277        Executor::SingleThread => {
278            for ia in 0..m.divceil(ker.mr()) {
279                for ib in 0..n.divceil(ker.nr()) {
280                    scratch.run(ker, non_linear, ia, ib)?;
281                }
282            }
283            Ok(())
284        }
285        #[cfg(feature = "multithread-mm")]
286        Executor::MultiThread(pool) => pool.install(|| {
287            pool.install(|| {
288                (0..m.div_ceil(ker.mr())).into_par_iter().try_for_each(|ia| {
289                    for ib in 0..n.divceil(ker.nr()) {
290                        scratch.run(ker, non_linear, ia, ib)?;
291                    }
292                    Ok(())
293                })
294            })
295        }),
296    }
297}
tract_linalg/frame/mmm/mod.rs

tract_linalg/frame/mmm/
mod.rs
