tract_core/ops/matmul/
pack.rs

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

tract_core/ops/matmul/
pack.rs
