tract_linalg/frame/mmm/
kernel.rs

1use crate::frame::pack::PackedFormat;
2
3use super::*;
4use std::borrow::Cow;
5use std::fmt::Debug;
6
7use crate::LADatum;
8
9pub trait MatMatMulKer: Clone + Debug + Send + Sync + 'static {
10    type Acc: LADatum;
11    fn name(&self) -> &str;
12    fn kernel(&self, op: &[FusedKerSpec<Self::Acc>]) -> isize;
13    fn mr(&self) -> usize;
14    fn nr(&self) -> usize;
15
16    fn quality(&self) -> ImplementationQuality;
17    fn dynamic_boost(&self) -> isize;
18
19    #[allow(clippy::type_complexity)]
20    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)];
21    fn stores(&self) -> Cow<'_, [DatumType]>;
22
23    #[allow(unused_variables)]
24    fn can_fuse(&self, spec: &FusedSpec) -> bool {
25        true
26    }
27
28    #[allow(unused_variables)]
29    fn is_supported_here(&self) -> bool {
30        true
31    }
32
33    /// Whether the border-tile store scratch should be laid out row-major
34    /// (n contiguous) instead of the default column-major (mr contiguous).
35    /// Set by kernels whose store has an aligned row-major bulk path.
36    fn stores_row_major_tile(&self) -> bool {
37        false
38    }
39}
40
41type Kernel<Acc> = unsafe fn(&[FusedKerSpec<Acc>]) -> isize;
42
43#[derive(Clone)]
44pub struct DynKernel<const MR: usize, const NR: usize, Acc: LADatum> {
45    pub name: String,
46    pub kernel: Kernel<Acc>,
47    pub quality: ImplementationQuality,
48    pub packings: Vec<(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)>,
49    pub stores: Vec<DatumType>,
50    pub supported_predicate: fn() -> bool,
51    pub boost: fn() -> isize,
52    pub can_fuse: fn(&FusedSpec) -> bool,
53    pub row_major_store: bool,
54}
55
56impl<const MR: usize, const NR: usize, Acc: LADatum> DynKernel<MR, NR, Acc> {
57    pub fn new(
58        name: &str,
59        kernel: Kernel<Acc>,
60        packing_a: PackedFormat,
61        packing_b: PackedFormat,
62        quality: ImplementationQuality,
63    ) -> Self {
64        let kernel = DynKernel {
65            name: name.to_string(),
66            kernel,
67            quality,
68            packings: vec![],
69            stores: vec![Acc::datum_type()],
70            supported_predicate: || true,
71            boost: || 0,
72            can_fuse: |_| true,
73            row_major_store: false,
74        };
75        kernel.with_packing(packing_a, packing_b)
76    }
77
78    pub fn with_platform_condition(mut self, f: fn() -> bool) -> Self {
79        self.supported_predicate = f;
80        self
81    }
82
83    pub fn with_boost(mut self, f: fn() -> isize) -> Self {
84        self.boost = f;
85        self
86    }
87
88    pub fn with_packing(mut self, a: impl MMMInputFormat, b: impl MMMInputFormat) -> Self {
89        self.packings.push((Box::new(a), Box::new(b)));
90        self
91    }
92
93    pub fn with_packing_a(self, a: impl MMMInputFormat) -> Self {
94        let b = self.regular_pack_b();
95        self.with_packing(a, b)
96    }
97
98    pub fn regular_pack_a(&self) -> PackedFormat {
99        *self.packings[0].0.clone().downcast::<PackedFormat>().unwrap()
100    }
101
102    pub fn regular_pack_b(&self) -> PackedFormat {
103        *self.packings[0].1.clone().downcast::<PackedFormat>().unwrap()
104    }
105
106    pub fn with_can_fuse(self, can_fuse: fn(&FusedSpec) -> bool) -> Self {
107        Self { can_fuse, ..self }
108    }
109
110    pub fn with_store<D: LADatum>(mut self) -> Self {
111        self.stores.push(D::datum_type());
112        self
113    }
114
115    pub fn mmm(&self) -> Box<dyn MatMatMul> {
116        Box::new(self.clone())
117    }
118}
119
120impl<const MR: usize, const NR: usize, Acc: LADatum> Debug for DynKernel<MR, NR, Acc> {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        write!(f, "{}", self.name)
123    }
124}
125
126impl<const MR: usize, const NR: usize, Acc: LADatum> MatMatMulKer for DynKernel<MR, NR, Acc> {
127    type Acc = Acc;
128    fn name(&self) -> &str {
129        &self.name
130    }
131
132    fn mr(&self) -> usize {
133        MR
134    }
135
136    fn nr(&self) -> usize {
137        NR
138    }
139
140    fn quality(&self) -> ImplementationQuality {
141        self.quality
142    }
143
144    fn is_supported_here(&self) -> bool {
145        (self.supported_predicate)()
146    }
147
148    fn can_fuse(&self, spec: &FusedSpec) -> bool {
149        (self.can_fuse)(spec)
150    }
151
152    fn kernel(&self, op: &[FusedKerSpec<Self::Acc>]) -> isize {
153        unsafe { (self.kernel)(op) }
154    }
155
156    #[allow(clippy::type_complexity)]
157    fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
158        &self.packings
159    }
160
161    fn stores(&self) -> Cow<'_, [DatumType]> {
162        Cow::Borrowed(&self.stores)
163    }
164
165    fn dynamic_boost(&self) -> isize {
166        (self.boost)()
167    }
168
169    fn stores_row_major_tile(&self) -> bool {
170        self.row_major_store
171    }
172}
tract_linalg/frame/mmm/kernel.rs

tract_linalg/frame/mmm/
kernel.rs
