tract_core/ops/einsum/
prefix_matmul.rs

1use tract_data::itertools::Itertools;
2use tract_linalg::Scaler;
3use tract_ndarray::Ix2;
4use tract_num_traits::One;
5
6use super::einsum_matmul::EinSumMatMul;
7use super::eval::dequant_inputs;
8use crate::internal::*;
9use crate::ops::einsum::block_quant_aware_input_shape;
10use crate::ops::konst::Const;
11
12#[derive(Debug, Default)]
13pub struct EinSumToPrefixMatmulCtx {
14    pub ensure_strict_matmul_semantic: bool,
15}
16
17pub fn rewrite_einsum_to_prefix_matmul(
18    model: &mut TypedModel,
19    ensure_strict_matmul_semantic: bool,
20) -> TractResult<()> {
21    super::einsum_matmul::merge_consecutive_same_role_axes(model)?;
22    super::einsum_matmul::detect_all(model)?;
23    let ctx = EinSumToPrefixMatmulCtx { ensure_strict_matmul_semantic };
24    Rewriter::default().with_rule_for("einsum-to-prefix-matmul", rule).rewrite(&ctx, model)
25}
26
27fn rule(
28    ctx: &EinSumToPrefixMatmulCtx,
29    model: &TypedModel,
30    node: &TypedNode,
31    node_name: &str,
32    op: &EinSumMatMul,
33) -> TractResult<Option<TypedModelPatch>> {
34    // F: 2 inputs
35    // Q: 9 inputs
36    let is_fp_mm = op.q_params.is_none() && node.inputs.len() == 2;
37    let is_q_mm = op.q_params.is_some() && node.inputs.len() == 9;
38    rule_if!(is_fp_mm || is_q_mm);
39    // PrefixMatMul carries no bias; the quantized fusion would drop the einsum
40    // bias input. Only the serialization path (strict) relies on this form, so
41    // on exec backends leave quantized einsums to lower through requant.
42    rule_if!(ctx.ensure_strict_matmul_semantic || op.q_params.is_none());
43    rule_if!(
44        op.q_params.is_none()
45            || model.node_input_facts(node.id)?.iter().skip(3).all(|i| i.konst.is_some())
46    );
47    let prefix: String = op
48        .axes
49        .iter_all_axes()
50        .filter(|a| ![op.m_axis, op.k_axis, op.n_axis].contains(&a.repr))
51        .map(|a| a.repr)
52        .collect();
53    let mut patch = TypedModelPatch::default();
54    let inputs = patch.taps(model, &node.inputs)?;
55    let mut wire = tvec!(inputs[0], inputs[1]);
56
57    let (m, k, n) = (op.m_axis, op.k_axis, op.n_axis);
58    let a_order_es: String = op.axes.axes(InOut::In(0)).map(|a| a.repr).collect();
59    let a_order_mm = format!("{prefix}{m}{k}");
60    let a_order_mm_t = format!("{prefix}{k}{m}");
61    let a_transform =
62        format!("{a_order_es}->{a_order_mm}").parse::<AxesMapping>()?.translate_to_axis_ops()?;
63    let a_transform_t =
64        format!("{a_order_es}->{a_order_mm_t}").parse::<AxesMapping>()?.translate_to_axis_ops()?;
65    let transpose_a = a_transform.len() > a_transform_t.len();
66    let a_transform = if transpose_a { a_transform_t } else { a_transform };
67    let name = format!("{node_name}.fix_a");
68    for op in a_transform {
69        wire[0] = patch.wire_node(&name, op, &[wire[0]])?[0];
70    }
71    // terrible hack to maintain exotic fact through eager propagation of constant through the
72    // axes transformation
73    if let Some(op) = patch.node_mut(wire[0].node).op_as_mut::<Const>() {
74        *op = Const::new_with_opt_exotic_fact(
75            op.val().clone(),
76            model.outlet_fact(node.inputs[0])?.exotic_fact.clone(),
77        )?;
78    }
79    patch
80        .outlet_fact_mut(wire[0])?
81        .exotic_fact
82        .clone_from(&model.outlet_fact(node.inputs[0])?.exotic_fact);
83    // end of hack
84
85    let b_order_es: String = op.axes.axes(InOut::In(1)).map(|a| a.repr).collect();
86    let b_order_mm = format!("{prefix}{k}{n}");
87    let b_order_mm_t = format!("{prefix}{n}{k}");
88    let b_transform =
89        format!("{b_order_es}->{b_order_mm}").parse::<AxesMapping>()?.translate_to_axis_ops()?;
90    let b_transform_t =
91        format!("{b_order_es}->{b_order_mm_t}").parse::<AxesMapping>()?.translate_to_axis_ops()?;
92    let transpose_b = b_transform.len() > b_transform_t.len();
93    let b_transform = if transpose_b { b_transform_t } else { b_transform };
94    let name = format!("{node_name}.fix_b");
95    for op in b_transform {
96        wire[1] = patch.wire_node(&name, op, &[wire[1]])?[0];
97    }
98
99    let c_order_es: String = op.axes.axes(InOut::Out(0)).map(|a| a.repr).collect();
100    let c_order_mm = format!("{prefix}{m}{n}");
101    let c_order_mm_t = format!("{prefix}{n}{m}");
102    let c_transform =
103        format!("{c_order_mm}->{c_order_es}").parse::<AxesMapping>()?.translate_to_axis_ops()?;
104    let c_transform_t =
105        format!("{c_order_mm_t}->{c_order_es}").parse::<AxesMapping>()?.translate_to_axis_ops()?;
106    let transpose_c = c_transform.len() > c_transform_t.len();
107    let c_transform = if transpose_c { c_transform_t } else { c_transform };
108    let quantize_output = if let Some(qp) = op.q_params {
109        let qparams: Vec<&Tensor> = inputs[3..9]
110            .iter()
111            .map(|f| {
112                patch
113                    .outlet_fact(*f)?
114                    .konst
115                    .as_deref()
116                    .context("Can only translate fixed scalar quantization")
117            })
118            .try_collect()?;
119        Some(qp.with_qparams(QParams::ZpScale {
120            zero_point: qparams[4].cast_to_scalar::<i32>()?,
121            scale: qparams[5].cast_to_scalar::<f32>()?,
122        }))
123    } else {
124        None
125    };
126
127    let operating_dt = if ctx.ensure_strict_matmul_semantic {
128        let input_facts = model.node_input_facts(node.id)?;
129        let a_dt = input_facts[0].datum_type;
130        let b_dt = input_facts[1].datum_type;
131        let operating_dt = quantize_output.unwrap_or(op.operating_dt);
132        let a_plain = input_facts[0].is_plain();
133        let b_plain = input_facts[1].is_plain();
134        let allowed_dt = matmul_semantic_output_dt(&a_dt, a_plain, &b_dt, b_plain);
135
136        ensure!(
137            operating_dt == allowed_dt,
138            format!(
139                "Strict matmul semantic require operating_dt to be {allowed_dt:?} \
140                for (a: {a_dt:?}, b:{b_dt:?}) but got {:?}.",
141                op.operating_dt
142            )
143        );
144
145        None
146    } else {
147        Some(op.operating_dt)
148    };
149
150    wire = patch.wire_node(
151        node_name,
152        PrefixMatMul { transpose_a, transpose_b, transpose_c, quantize_output, operating_dt },
153        &wire,
154    )?;
155
156    for (ix, op) in c_transform.into_iter().enumerate() {
157        wire = patch.wire_node(format!("{node_name}.fix_c.{ix}"), op, &wire)?;
158    }
159    patch.shunt_outside(model, node.id.into(), wire[0])?;
160    Ok(Some(patch))
161}
162
163fn matmul_semantic_output_dt(
164    a_dt: &DatumType,
165    a_plain: bool,
166    b_dt: &DatumType,
167    b_plain: bool,
168) -> DatumType {
169    if a_plain && a_dt.is_number() {
170        *a_dt
171    } else if b_plain && b_dt.is_number() {
172        *b_dt
173    } else if a_dt.is_number() {
174        *a_dt
175    } else if b_dt.is_number() {
176        *b_dt
177    } else {
178        f32::datum_type()
179    }
180}
181
182#[derive(Clone, Debug, Copy, PartialEq, Eq)]
183pub struct PrefixMatMul {
184    pub transpose_a: bool,
185    pub transpose_b: bool,
186    pub transpose_c: bool,
187    pub quantize_output: Option<DatumType>,
188    pub operating_dt: Option<DatumType>,
189}
190
191impl PrefixMatMul {
192    fn output_shape<D: DimLike + One>(&self, a: &[D], b: &[D]) -> TVec<D> {
193        let rank = a.len();
194        let mut output: TVec<D> = (0..rank - 2)
195            .map(|ix| if a[ix].is_one() { b[ix].clone() } else { a[ix].clone() })
196            .collect();
197        output.push(a[rank - 2 + self.transpose_a as usize].clone());
198        output.push(b[rank - 2 + !self.transpose_b as usize].clone());
199        if self.transpose_c {
200            output.swap(rank - 2, rank - 1);
201        }
202        output
203    }
204
205    fn mm<Acc: Datum + tract_ndarray::LinalgScalar>(
206        &self,
207        acc: &mut Tensor,
208        a: &Tensor,
209        b: &Tensor,
210    ) -> TractResult<()> {
211        use crate::ndarray::Dimension;
212        let casted_a = a.cast_to::<Acc>()?;
213        let a = casted_a.to_plain_array_view::<Acc>()?;
214        let casted_b = b.cast_to::<Acc>()?;
215        let b = casted_b.to_plain_array_view::<Acc>()?;
216        let mut c_plain = acc.try_as_plain_mut()?;
217        let mut c = c_plain.to_array_view_mut::<Acc>()?;
218        for prefix in tract_ndarray::indices(&c.shape()[..c.ndim() - 2]) {
219            let mut a = a.view();
220            let mut b = b.view();
221            let mut c = c.view_mut();
222            for &d in prefix.slice().iter() {
223                a.index_axis_inplace(tract_ndarray::Axis(0), d.min(a.shape()[0] - 1));
224                b.index_axis_inplace(tract_ndarray::Axis(0), d.min(b.shape()[0] - 1));
225                c.index_axis_inplace(tract_ndarray::Axis(0), d);
226            }
227            let a = a.into_dimensionality::<Ix2>().unwrap();
228            let b = b.into_dimensionality::<Ix2>().unwrap();
229            let mut c = c.into_dimensionality::<Ix2>().unwrap();
230            let a = if self.transpose_a { a.t() } else { a };
231            let b = if self.transpose_b { b.t() } else { b };
232            if self.transpose_c { c.assign(&b.t().dot(&a.t())) } else { c.assign(&a.dot(&b)) }
233        }
234        Ok(())
235    }
236}
237
238impl Op for PrefixMatMul {
239    fn name(&self) -> StaticName {
240        "PrefixMatMul".into()
241    }
242
243    fn info(&self) -> TractResult<Vec<String>> {
244        Ok(vec![format!(
245            "transpose_a: {} transpose_b: {} transpose_c: {} q: {:?}",
246            self.transpose_a, self.transpose_b, self.transpose_c, self.quantize_output
247        )])
248    }
249
250    op_as_typed_op!();
251}
252
253impl EvalOp for PrefixMatMul {
254    fn is_stateless(&self) -> bool {
255        true
256    }
257
258    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
259        let c_dt = self.operating_dt.unwrap_or_else(|| {
260            let a_dt = inputs[0].datum_type();
261            let b_dt = inputs[1].datum_type();
262            matmul_semantic_output_dt(&a_dt, inputs[0].is_plain(), &b_dt, inputs[1].is_plain())
263        });
264
265        let inputs = dequant_inputs(c_dt, inputs)?;
266
267        let output_shape = self.output_shape(inputs[0].shape(), inputs[1].shape());
268
269        if let Some(qp) = self.quantize_output {
270            let mut acc = Tensor::zero_dt(i32::datum_type(), &output_shape)?;
271            let mut a_i32 = inputs[0].cast_to::<i32>()?.into_owned();
272            a_i32
273                .try_as_plain_mut()?
274                .as_slice_mut::<i32>()?
275                .iter_mut()
276                .for_each(|x| *x -= inputs[0].datum_type().zp_scale().0);
277            let mut b_i32 = inputs[1].cast_to::<i32>()?.into_owned();
278            b_i32
279                .try_as_plain_mut()?
280                .as_slice_mut::<i32>()?
281                .iter_mut()
282                .for_each(|x| *x -= inputs[1].datum_type().zp_scale().0);
283            self.mm::<i32>(&mut acc, &a_i32, &b_i32)?;
284            let scale = inputs[0].datum_type().zp_scale().1 * inputs[1].datum_type().zp_scale().1
285                / qp.zp_scale().1;
286            let scaler = Scaler::new(scale, tract_linalg::mmm::RoundingPolicy::Even);
287            acc.to_plain_array_view_mut::<i32>()?.iter_mut().for_each(|x| *x = *x * scaler);
288            let mut c: Tensor = acc.cast_to_dt(qp.unquantized())?.into_owned();
289            unsafe { c.set_datum_type(qp) };
290            Ok(tvec!(c.into_tvalue()))
291        } else {
292            let mut c = Tensor::zero_dt(c_dt, &output_shape)?;
293            dispatch_floatlike!(Self::mm(c_dt)(self, &mut c, &inputs[0], &inputs[1]))?;
294            Ok(tvec!(c.into_tvalue()))
295        }
296    }
297}
298
299impl TypedOp for PrefixMatMul {
300    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
301        let [a, b] = inputs else {
302            bail!("Expects 2 inputs");
303        };
304        let a_shape = block_quant_aware_input_shape(a)?;
305        let b_shape = block_quant_aware_input_shape(b)?;
306        let dt = self.quantize_output.or(self.operating_dt).unwrap_or(matmul_semantic_output_dt(
307            &a.datum_type,
308            a.is_plain(),
309            &b.datum_type,
310            b.is_plain(),
311        ));
312        Ok(tvec!(dt.fact(self.output_shape(&a_shape, &b_shape))))
313    }
314
315    as_op!();
316}
317
318#[cfg(test)]
319mod test {
320    use crate::ops::einsum::EinSum;
321
322    use super::*;
323    use proptest::collection::vec;
324    use proptest::prelude::*;
325    use proptest::test_runner::{TestCaseResult, TestRunner};
326    use tract_data::itertools::Itertools;
327
328    pub fn tensor(shape: &[usize]) -> BoxedStrategy<Tensor> {
329        let shape = shape.to_vec();
330        let len = shape.iter().product::<usize>();
331        vec((-10i8..=10i8).prop_map(|i| i as f32), len..=len)
332            .prop_map(move |vec| tensor1(&vec).into_shape(&shape).unwrap())
333            .boxed()
334    }
335
336    fn full_shapes(e: &AxesMapping) -> BoxedStrategy<(Vec<usize>, Vec<usize>)> {
337        let e = e.clone();
338        let inputs_axes = e
339            .iter_all_axes()
340            .filter(|axis| axis.inputs[0].len() + axis.inputs[1].len() > 0)
341            .cloned()
342            .collect_vec();
343        let dims = vec![2usize..6; inputs_axes.len()];
344        dims.prop_map(move |dims| {
345            let a: Vec<usize> = e
346                .axes(InOut::In(0))
347                .map(|a| dims[inputs_axes.iter().position(|b| a == b).unwrap()])
348                .collect_vec();
349            let b: Vec<usize> = e
350                .axes(InOut::In(1))
351                .map(|a| dims[inputs_axes.iter().position(|b| a == b).unwrap()])
352                .collect_vec();
353            (a, b)
354        })
355        .boxed()
356    }
357
358    fn test_expr(expr: &str) -> TestCaseResult {
359        let expr = expr.to_string();
360        let mut runner = TestRunner::default();
361        let axes: AxesMapping = expr.parse().unwrap();
362        fn is_k(axes: &AxesMapping, input: usize, position: usize) -> bool {
363            let axis = axes.axis((InOut::In(input), position)).unwrap();
364            axis.inputs[1 - input].len() == 1 && axis.outputs[0].len() == 0
365        }
366        fn is_disapearing_axis(axes: &AxesMapping, input: usize, position: usize) -> bool {
367            let axis = axes.axis((InOut::In(input), position)).unwrap();
368            axis.outputs[0].len() == 0
369        }
370        let cases = full_shapes(&axes)
371            .prop_flat_map(|(a, b)| {
372                (
373                    a.iter()
374                        .enumerate()
375                        .map(|(ix, d)| {
376                            if is_k(&axes, 0, ix) {
377                                prop_oneof![Just(*d)].boxed()
378                            } else if is_disapearing_axis(&axes, 0, ix) {
379                                Just(1).boxed()
380                            } else {
381                                prop_oneof![Just(1usize), Just(*d)].boxed()
382                            }
383                        })
384                        .collect_vec(),
385                    b.iter()
386                        .enumerate()
387                        .map(|(ix, d)| {
388                            if is_k(&axes, 1, ix) {
389                                prop_oneof![Just(*d)].boxed()
390                            } else if is_disapearing_axis(&axes, 1, ix) {
391                                Just(1).boxed()
392                            } else {
393                                prop_oneof![Just(1usize), Just(*d)].boxed()
394                            }
395                        })
396                        .collect_vec(),
397                )
398            })
399            .prop_flat_map(|(a_shape, b_shape)| (tensor(&a_shape), tensor(&b_shape)))
400            .prop_map(|(a, b)| EinSumProblem { expr: expr.clone(), a, b });
401        runner.run(&cases, |pb| pb.check().map_err(|e| TestCaseError::fail(e.to_string())))?;
402        Ok(())
403    }
404
405    #[derive(Debug, Clone, PartialEq, Eq)]
406    struct EinSumProblem {
407        expr: String,
408        a: Tensor,
409        b: Tensor,
410    }
411
412    impl EinSumProblem {
413        fn check(&self) -> TractResult<()> {
414            let mut model = TypedModel::default();
415            let sa = model.add_source("a", f32::fact(self.a.shape()))?;
416            let sb = model.add_source("b", f32::fact(self.b.shape()))?;
417            let einsum = model.wire_node(
418                "einsum",
419                EinSum::new(self.expr.parse().unwrap(), f32::datum_type()),
420                &[sa, sb],
421            )?;
422            model.select_output_outlets(&einsum)?;
423            let a = self.a.clone().into_tvalue();
424            let b = self.b.clone().into_tvalue();
425            let inputs = tvec!(a, b);
426            let reference = TypedRunnableModel::new(model.clone())?.run(inputs.clone())?.remove(0);
427            rewrite_einsum_to_prefix_matmul(&mut model, true)?;
428            assert!(model.nodes.iter().all(|n| !n.op_is::<EinSum>()));
429            let test = TypedRunnableModel::new(model)?.run(inputs)?.remove(0);
430            reference.close_enough(&test, true)
431        }
432    }
433
434    #[rustfmt::skip] #[test] fn prop_mk_kn_mn() -> TestCaseResult { test_expr("mk,kn->mn") }
435    #[rustfmt::skip] #[test] fn prop_km_kn_mn() -> TestCaseResult { test_expr("km,kn->mn") }
436    #[rustfmt::skip] #[test] fn prop_mk_nk_mn() -> TestCaseResult { test_expr("mk,nk->mn") }
437    #[rustfmt::skip] #[test] fn prop_mk_kn_nm() -> TestCaseResult { test_expr("mk,kn->nm") }
438    #[rustfmt::skip] #[test] fn prop_k_kn_mn() -> TestCaseResult { test_expr("k,kn->mn") }
439    #[rustfmt::skip] #[test] fn prop_mk_k_mn() -> TestCaseResult { test_expr("mk,k->mn") }
440    #[rustfmt::skip] #[test] fn prop_m_n_mn() -> TestCaseResult { test_expr("m,n->mn") }
441    #[rustfmt::skip] #[test] fn prop_amk_akn_amn() -> TestCaseResult { test_expr("amk,akn->amn") }
442    #[rustfmt::skip] #[test] fn prop_mk_akn_amn() -> TestCaseResult { test_expr("mk,akn->amn") }
443    #[rustfmt::skip] #[test] fn prop_btgi_gih_tgh() -> TestCaseResult { test_expr("btgi,gih->tgh") }
444    #[rustfmt::skip] #[test] fn prop_tgi_gih_btgh() -> TestCaseResult { test_expr("tgi,gih->btgh") }
445
446    #[test]
447    fn k_kn_mn_0() -> TractResult<()> {
448        EinSumProblem {
449            expr: "k,kn->mn".to_string(),
450            a: tensor1(&[0f32, 0f32]),
451            b: tensor2(&[[0f32, 0.], [0., 0.]]),
452        }
453        .check()
454    }
455
456    #[test]
457    fn mk_k_mn_0() -> TractResult<()> {
458        EinSumProblem {
459            expr: "mk,k->mn".to_string(),
460            a: Tensor::zero::<f32>(&[2, 2]).unwrap(),
461            b: Tensor::zero::<f32>(&[2]).unwrap(),
462        }
463        .check()
464    }
465
466    #[test]
467    fn mk_k_mn_1() -> TractResult<()> {
468        EinSumProblem {
469            expr: "mk,k->mn".to_string(),
470            a: Tensor::zero::<f32>(&[1, 2]).unwrap(),
471            b: Tensor::zero::<f32>(&[2]).unwrap(),
472        }
473        .check()
474    }
475
476    #[test]
477    fn mk_kn_nm_0() -> TractResult<()> {
478        EinSumProblem {
479            expr: "mk,kn->mn".to_string(),
480            a: Tensor::zero::<f32>(&[3, 2]).unwrap(),
481            b: Tensor::zero::<f32>(&[2, 2]).unwrap(),
482        }
483        .check()
484    }
485
486    #[test]
487    fn amk_akn_amn_0() -> TractResult<()> {
488        EinSumProblem {
489            expr: "amk,akn->amn".to_string(),
490            a: Tensor::zero::<f32>(&[1, 1, 2]).unwrap(),
491            b: Tensor::zero::<f32>(&[1, 2, 1]).unwrap(),
492        }
493        .check()
494    }
495
496    #[test]
497    fn amk_akn_amn_1() -> TractResult<()> {
498        EinSumProblem {
499            expr: "amk,akn->amn".to_string(),
500            a: Tensor::zero::<f32>(&[2, 1, 2]).unwrap(),
501            b: Tensor::zero::<f32>(&[1, 2, 1]).unwrap(),
502        }
503        .check()
504    }
505
506    #[test]
507    fn amk_akn_amn_2() -> TractResult<()> {
508        EinSumProblem {
509            expr: "amk,akn->amn".to_string(),
510            a: Tensor::zero::<f32>(&[1, 1, 2]).unwrap(),
511            b: Tensor::zero::<f32>(&[2, 2, 2]).unwrap(),
512        }
513        .check()
514    }
515
516    #[test]
517    fn amk_akn_amn_3() -> TractResult<()> {
518        EinSumProblem {
519            expr: "amk,akn->amn".to_string(),
520            a: Tensor::zero::<f32>(&[1, 1, 2]).unwrap(),
521            b: Tensor::zero::<f32>(&[2, 2, 1]).unwrap(),
522        }
523        .check()
524    }
525
526    #[test]
527    fn km_anbck_bmn_0() -> TractResult<()> {
528        EinSumProblem {
529            expr: "km,anbck->bmn".to_string(),
530            a: Tensor::zero::<f32>(&[2, 1]).unwrap(),
531            b: Tensor::zero::<f32>(&[1, 1, 1, 1, 2]).unwrap(),
532        }
533        .check()
534    }
535
536    fn check_k1(expr: &str, a_shape: &[usize], b_shape: &[usize]) -> TractResult<()> {
537        let a_len = a_shape.iter().product::<usize>();
538        let b_len = b_shape.iter().product::<usize>();
539        let a =
540            tensor1(&(0..a_len).map(|i| (i + 1) as f32).collect_vec()).into_shape(a_shape).unwrap();
541        let b = tensor1(&(0..b_len).map(|i| (i + 7) as f32 * 0.5).collect_vec())
542            .into_shape(b_shape)
543            .unwrap();
544        EinSumProblem { expr: expr.to_string(), a, b }.check()
545    }
546
547    // K=1 means the einsum's contraction degenerates to broadcast-mul. detect_rule
548    // short-circuits to a Mul op rather than dispatching the GEMM kernel for a single
549    // FMA per tile. The gmk,Ngnk->Ngmn pattern with k=1 is what depthwise ConvTranspose
550    // lowers to, and the per-tile GEMM overhead dominates in that case.
551    #[test]
552    fn k1_amk_akn_amn() -> TractResult<()> {
553        check_k1("amk,akn->amn", &[2, 3, 1], &[2, 1, 4])
554    }
555
556    #[test]
557    fn k1_gmk_ngnk_ngmn() -> TractResult<()> {
558        check_k1("gmk,Ngnk->Ngmn", &[3, 2, 1], &[1, 3, 4, 1])
559    }
560
561    #[test]
562    fn k1_mk_kn_mn() -> TractResult<()> {
563        check_k1("mk,kn->mn", &[2, 1], &[1, 3])
564    }
565
566    // The unit_k_to_broadcast_mul declutter rule is scoped to einsums whose output is
567    // consumed by DeconvSum (the original ConvTranspose-K=1 target). For these isolated
568    // einsum tests the rule won't fire — outputs still go through OptMatMul correctly,
569    // just slower than a broadcast Mul would be. The end-to-end ConvTranspose case is
570    // covered by integration tests (DFN3 erb_dec, GTCRN) and verified bit-exact in the
571    // PR description.
572    #[test]
573    fn k1_no_k_axis_outer() -> TractResult<()> {
574        check_k1("m,n->mn", &[3], &[4])
575    }
576
577    #[test]
578    fn k1_high_rank_no_decov_sum() -> TractResult<()> {
579        // From a CI proptest failure: a high-rank einsum with K=1, no DeconvSum
580        // downstream → rule must not fire, OptMatMul handles it correctly.
581        check_k1("wmexk,wxnk->ewnxm", &[2, 1, 2, 2, 1], &[2, 2, 1, 1])
582    }
583
584    #[test]
585    fn k1_sdpa_shape_no_deconv_sum() -> TractResult<()> {
586        // Regression for SDPA failure mode: when head_dim=1, SDPA's score einsum
587        // looks like a K=1 case structurally, but it's followed by Softmax (not
588        // DeconvSum). Rule must NOT fire — Metal's SDPA fusion would otherwise break.
589        check_k1("bhmk,bhnk->bhmn", &[1, 3, 4, 1], &[1, 3, 4, 1])
590    }
591
592    #[test]
593    fn q() -> TractResult<()> {
594        let qp = QParams::ZpScale { zero_point: 0, scale: 0.1 };
595        let op = EinSum {
596            axes: "mk,kn,m,,,,,,->mn".parse()?,
597            operating_dt: i32::datum_type(),
598            q_params: Some(DatumType::QI8(qp)),
599        };
600        let mut model = TypedModelPatch::default();
601        let inputs = [
602            model.add_source("a", DatumType::QI8(qp).fact([3, 2]))?,
603            model.add_source("b", DatumType::QI8(qp).fact([2, 4]))?,
604            model.add_source("bias", i32::datum_type().fact([3]))?,
605            model.add_const("a0", tensor0(qp.zp_scale().0))?,
606            model.add_const("a_scale", tensor0(qp.zp_scale().1))?,
607            model.add_const("b0", tensor0(qp.zp_scale().0))?,
608            model.add_const("b_scale", tensor0(qp.zp_scale().1))?,
609            model.add_const("c0", tensor0(qp.zp_scale().0))?,
610            model.add_const("c_scale", tensor0(qp.zp_scale().1))?,
611        ];
612        let wire = model.wire_node("einsum", op.clone(), &inputs)?;
613        model.select_output_outlets(&wire)?;
614        rewrite_einsum_to_prefix_matmul(&mut model, true)?;
615        assert!(model.nodes.iter().all(|n| !n.op_is::<EinSum>()));
616        Ok(())
617    }
618}
tract_core/ops/einsum/prefix_matmul.rs

tract_core/ops/einsum/
prefix_matmul.rs
