tract_core/ops/matmul/
pack.rs

1use crate::axes::Axis;
2use crate::internal::*;
3use ndarray::*;
4use tract_linalg::block_quant::{
5    BlockQuantStorage, PackedBlockQuantFact, PackedBlockQuantFormat, block_quant_slice,
6};
7use tract_linalg::mmm::{MMMInputFormat, MMMInputValue, PackedMatrixStorage};
8
9use super::ModePicker;
10
11#[derive(Debug, Clone, PartialEq, Eq, Hash)]
12pub struct OptMatMulPack {
13    pub(crate) packers: Vec<Box<dyn MMMInputFormat>>,
14    pub(crate) mode_picker: ModePicker,
15    pub(crate) k_axis: usize,
16    pub(crate) mn_axis: usize,
17}
18
19impl Op for OptMatMulPack {
20    fn name(&self) -> StaticName {
21        "OptMatMulPack".into()
22    }
23
24    fn info(&self) -> TractResult<Vec<String>> {
25        Ok(vec![format!("{:?}. k axis: {}, mn axis: {}", self.packers, self.k_axis, self.mn_axis)])
26    }
27
28    op_as_typed_op!();
29}
30
31impl EvalOp for OptMatMulPack {
32    op_out_of_plan!();
33
34    fn eval(&self, ctx: &EvalContext, mut inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
35        self.do_eval(ctx, inputs.remove(0))
36    }
37}
38
39impl TypedOp for OptMatMulPack {
40    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
41        match self.mode_picker {
42            ModePicker::Single => ensure!(self.packers.len() == 1),
43            ModePicker::VecVsMat => ensure!(self.packers.len() == 2),
44        }
45        let k = inputs[0].shape[self.k_axis].clone();
46        let mn = inputs[0].shape[self.mn_axis].clone();
47        let exotic_fact = DynPackedExoticFact { k, mn, packers: self.packers.clone() };
48        Ok(tvec!(
49            inputs[0]
50                .datum_type
51                .fact(self.output_shape(&inputs[0].shape))
52                .with_exotic_fact(exotic_fact)
53        ))
54    }
55
56    fn axes_mapping(
57        &self,
58        inputs: &[&TypedFact],
59        outputs: &[&TypedFact],
60    ) -> TractResult<AxesMapping> {
61        let mut axes: Vec<Axis> = (0..inputs[0].rank())
62            .filter(|&ix| ix != self.k_axis && ix != self.mn_axis)
63            .enumerate()
64            .zip('a'..)
65            .map(|((o, i), repr)| Axis::new(repr, 1, 1).input(0, i).output(0, o))
66            .collect();
67        axes.push(Axis::new('K', 1, 1).input(0, self.k_axis));
68        axes.push(Axis::new('M', 1, 1).input(0, self.mn_axis));
69        axes.push(Axis::new('P', 1, 1).output(0, outputs[0].rank()));
70        AxesMapping::new(1, 1, axes)
71    }
72
73    as_op!();
74}
75
76impl OptMatMulPack {
77    fn do_eval(&self, _ctx: &EvalContext, input: TValue) -> TractResult<TVec<TValue>> {
78        let mode = self.mode_picker.pick(input.shape()[self.mn_axis])?;
79        let packer = &self.packers[mode];
80        let output_shape: TVec<usize> = self.output_shape(input.shape());
81        if let Some(strided) = packer.downcast_ref::<tract_linalg::strided_panel::StridedKMajor>() {
82            // Share the activation allocation. A copy here is the pack this format exists to skip.
83            let arc = input.as_arc_tensor().unwrap().clone();
84            let values = if output_shape.iter().all(|d| *d == 1) {
85                tvec![strided.panel_from_arc(arc, 0, self.k_axis, self.mn_axis)?]
86            } else {
87                let mut bc_shape: TVec<usize> = input.shape().into();
88                bc_shape[self.k_axis] = 1;
89                bc_shape[self.mn_axis] = 1;
90                let mut values: TVec<Box<dyn MMMInputValue>> =
91                    TVec::with_capacity(output_shape.iter().product());
92                for coord in indices(&*bc_shape) {
93                    let offset = coord
94                        .as_array_view()
95                        .iter()
96                        .zip(input.strides())
97                        .map(|(x, s)| *x as isize * s)
98                        .sum::<isize>()
99                        * input.datum_type().size_of() as isize;
100                    values.push(strided.panel_from_arc(
101                        arc.clone(),
102                        offset,
103                        self.k_axis,
104                        self.mn_axis,
105                    )?);
106                }
107                values
108            };
109            let stores = PackedMatrixStorage::new_batched(&output_shape, values)
110                .into_tensor(input.datum_type());
111            return Ok(tvec!(stores.into_tvalue()));
112        }
113        unsafe {
114            let stores = if output_shape.iter().all(|d| *d == 1) {
115                let packed = packer.prepare_one_view(&input.view(), self.k_axis, self.mn_axis)?;
116                PackedMatrixStorage::new_batched(&output_shape, tvec![packed])
117                    .into_tensor(input.datum_type())
118            } else {
119                let mut bc_shape: TVec<usize> = input.shape().into();
120                bc_shape[self.k_axis] = 1;
121                bc_shape[self.mn_axis] = 1;
122
123                let mut values: TVec<Box<dyn MMMInputValue>> =
124                    TVec::with_capacity(output_shape.iter().product());
125                for coord in indices(&*bc_shape) {
126                    let offset = coord
127                        .as_array_view()
128                        .iter()
129                        .zip(input.strides())
130                        .map(|(x, s)| *x as isize * s)
131                        .sum::<isize>()
132                        * input.datum_type().size_of() as isize;
133                    let view =
134                        TensorView::from_bytes(&input, offset, input.shape(), input.strides());
135                    values.push(packer.prepare_one_view(&view, self.k_axis, self.mn_axis)?);
136                }
137                PackedMatrixStorage::new_batched(&output_shape, values)
138                    .into_tensor(input.datum_type())
139            };
140            Ok(tvec!(stores.into_tvalue()))
141        }
142    }
143
144    pub fn output_shape<D: DimLike>(&self, input: &[D]) -> TVec<D> {
145        let mut packed_shape: TVec<D> = input.into();
146        packed_shape.remove(self.mn_axis.max(self.k_axis));
147        packed_shape.remove(self.mn_axis.min(self.k_axis));
148        packed_shape
149    }
150}
151
152#[derive(Hash, Clone, Debug, PartialEq, Eq)]
153pub struct DynPackedExoticFact {
154    pub k: TDim,
155    pub mn: TDim,
156    pub packers: Vec<Box<dyn MMMInputFormat>>,
157}
158
159impl ExoticFact for DynPackedExoticFact {
160    fn buffer_sizes(&self) -> TVec<TDim> {
161        tvec!(self.packers[0].mem_size(self.k.clone(), self.mn.clone()))
162    }
163}
164
165#[derive(Debug, Clone, Hash, Eq, PartialEq)]
166pub struct OptSimpleMatMulPack {
167    pub(crate) packed_format: PackedBlockQuantFormat,
168    pub(crate) k: usize,
169    pub(crate) m: usize,
170}
171
172impl Op for OptSimpleMatMulPack {
173    fn name(&self) -> StaticName {
174        "OptSimpleMatMulPack".into()
175    }
176    op_as_typed_op!();
177}
178
179impl EvalOp for OptSimpleMatMulPack {
180    op_out_of_plan!();
181
182    fn state(&self, _ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
183        Ok(None)
184    }
185
186    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
187        let input = args_1!(inputs);
188        let bqs = input.try_storage_as::<BlockQuantStorage>()?;
189        // Leading dims before the last 2 (M, K) are batch/group dims
190        let num_groups: usize = input.shape()[..input.rank().saturating_sub(2)].iter().product();
191        let m_per_group = input.shape()[input.rank() - 2];
192        let k = *input.shape().last().unwrap();
193        let values = (0..num_groups)
194            .map(|g| {
195                let slice = block_quant_slice(bqs.value(), bqs.format(), m_per_group, k, g);
196                let iv: Box<dyn MMMInputValue> = Box::new(self.packed_format.pack(slice, k)?);
197                Ok(iv)
198            })
199            .collect::<TractResult<TVec<_>>>()?;
200        let leading_shape = &input.shape()[..input.rank().saturating_sub(2)];
201        let output =
202            PackedMatrixStorage::new_batched(leading_shape, values).into_tensor(input.datum_type());
203        Ok(tvec!(output.into_tvalue()))
204    }
205}
206
207impl TypedOp for OptSimpleMatMulPack {
208    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
209        let input = inputs[0];
210        // Input shape is [G, M, K] — output removes M and K, keeping leading dims
211        let output_shape: TVec<TDim> = if input.rank() > 2 {
212            input.shape[..input.rank() - 2].to_vec().into()
213        } else {
214            tvec!()
215        };
216        let fact =
217            inputs[0].datum_type.fact(&*output_shape).with_exotic_fact(PackedBlockQuantFact {
218                format: self.packed_format.clone(),
219                shape: tvec!(self.m, self.k),
220            });
221        Ok(tvec!(fact))
222    }
223
224    as_op!();
225}
tract_core/ops/matmul/pack.rs

tract_core/ops/matmul/
pack.rs
