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    op_out_of_plan!();
255
256    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
257        let c_dt = self.operating_dt.unwrap_or_else(|| {
258            let a_dt = inputs[0].datum_type();
259            let b_dt = inputs[1].datum_type();
260            matmul_semantic_output_dt(&a_dt, inputs[0].is_plain(), &b_dt, inputs[1].is_plain())
261        });
262
263        let inputs = dequant_inputs(c_dt, inputs)?;
264
265        let output_shape = self.output_shape(inputs[0].shape(), inputs[1].shape());
266
267        if let Some(qp) = self.quantize_output {
268            let mut acc = Tensor::zero_dt(i32::datum_type(), &output_shape)?;
269            let mut a_i32 = inputs[0].cast_to::<i32>()?.into_owned();
270            a_i32
271                .try_as_plain_mut()?
272                .as_slice_mut::<i32>()?
273                .iter_mut()
274                .for_each(|x| *x -= inputs[0].datum_type().zp_scale().0);
275            let mut b_i32 = inputs[1].cast_to::<i32>()?.into_owned();
276            b_i32
277                .try_as_plain_mut()?
278                .as_slice_mut::<i32>()?
279                .iter_mut()
280                .for_each(|x| *x -= inputs[1].datum_type().zp_scale().0);
281            self.mm::<i32>(&mut acc, &a_i32, &b_i32)?;
282            let scale = inputs[0].datum_type().zp_scale().1 * inputs[1].datum_type().zp_scale().1
283                / qp.zp_scale().1;
284            let scaler = Scaler::new(scale, tract_linalg::mmm::RoundingPolicy::Even);
285            acc.to_plain_array_view_mut::<i32>()?.iter_mut().for_each(|x| *x = *x * scaler);
286            let mut c: Tensor = acc.cast_to_dt(qp.unquantized())?.into_owned();
287            unsafe { c.set_datum_type(qp) };
288            Ok(tvec!(c.into_tvalue()))
289        } else {
290            let mut c = Tensor::zero_dt(c_dt, &output_shape)?;
291            dispatch_floatlike!(Self::mm(c_dt)(self, &mut c, &inputs[0], &inputs[1]))?;
292            Ok(tvec!(c.into_tvalue()))
293        }
294    }
295}
296
297impl TypedOp for PrefixMatMul {
298    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
299        let [a, b] = inputs else {
300            bail!("Expects 2 inputs");
301        };
302        let a_shape = block_quant_aware_input_shape(a)?;
303        let b_shape = block_quant_aware_input_shape(b)?;
304        let dt = self.quantize_output.or(self.operating_dt).unwrap_or(matmul_semantic_output_dt(
305            &a.datum_type,
306            a.is_plain(),
307            &b.datum_type,
308            b.is_plain(),
309        ));
310        Ok(tvec!(dt.fact(self.output_shape(&a_shape, &b_shape))))
311    }
312
313    as_op!();
314}
315
316#[cfg(test)]
317mod test {
318    use crate::ops::einsum::EinSum;
319
320    use super::*;
321    use proptest::collection::vec;
322    use proptest::prelude::*;
323    use proptest::test_runner::{TestCaseResult, TestRunner};
324    use tract_data::itertools::Itertools;
325
326    pub fn tensor(shape: &[usize]) -> BoxedStrategy<Tensor> {
327        let shape = shape.to_vec();
328        let len = shape.iter().product::<usize>();
329        vec((-10i8..=10i8).prop_map(|i| i as f32), len..=len)
330            .prop_map(move |vec| tensor1(&vec).into_shape(&shape).unwrap())
331            .boxed()
332    }
333
334    fn full_shapes(e: &AxesMapping) -> BoxedStrategy<(Vec<usize>, Vec<usize>)> {
335        let e = e.clone();
336        let inputs_axes = e
337            .iter_all_axes()
338            .filter(|axis| axis.inputs[0].len() + axis.inputs[1].len() > 0)
339            .cloned()
340            .collect_vec();
341        let dims = vec![2usize..6; inputs_axes.len()];
342        dims.prop_map(move |dims| {
343            let a: Vec<usize> = e
344                .axes(InOut::In(0))
345                .map(|a| dims[inputs_axes.iter().position(|b| a == b).unwrap()])
346                .collect_vec();
347            let b: Vec<usize> = e
348                .axes(InOut::In(1))
349                .map(|a| dims[inputs_axes.iter().position(|b| a == b).unwrap()])
350                .collect_vec();
351            (a, b)
352        })
353        .boxed()
354    }
355
356    fn test_expr(expr: &str) -> TestCaseResult {
357        let expr = expr.to_string();
358        let mut runner = TestRunner::default();
359        let axes: AxesMapping = expr.parse().unwrap();
360        fn is_k(axes: &AxesMapping, input: usize, position: usize) -> bool {
361            let axis = axes.axis((InOut::In(input), position)).unwrap();
362            axis.inputs[1 - input].len() == 1 && axis.outputs[0].len() == 0
363        }
364        fn is_disapearing_axis(axes: &AxesMapping, input: usize, position: usize) -> bool {
365            let axis = axes.axis((InOut::In(input), position)).unwrap();
366            axis.outputs[0].len() == 0
367        }
368        let cases = full_shapes(&axes)
369            .prop_flat_map(|(a, b)| {
370                (
371                    a.iter()
372                        .enumerate()
373                        .map(|(ix, d)| {
374                            if is_k(&axes, 0, ix) {
375                                prop_oneof![Just(*d)].boxed()
376                            } else if is_disapearing_axis(&axes, 0, ix) {
377                                Just(1).boxed()
378                            } else {
379                                prop_oneof![Just(1usize), Just(*d)].boxed()
380                            }
381                        })
382                        .collect_vec(),
383                    b.iter()
384                        .enumerate()
385                        .map(|(ix, d)| {
386                            if is_k(&axes, 1, ix) {
387                                prop_oneof![Just(*d)].boxed()
388                            } else if is_disapearing_axis(&axes, 1, ix) {
389                                Just(1).boxed()
390                            } else {
391                                prop_oneof![Just(1usize), Just(*d)].boxed()
392                            }
393                        })
394                        .collect_vec(),
395                )
396            })
397            .prop_flat_map(|(a_shape, b_shape)| (tensor(&a_shape), tensor(&b_shape)))
398            .prop_map(|(a, b)| EinSumProblem { expr: expr.clone(), a, b });
399        runner.run(&cases, |pb| pb.check().map_err(|e| TestCaseError::fail(e.to_string())))?;
400        Ok(())
401    }
402
403    #[derive(Debug, Clone, PartialEq, Eq)]
404    struct EinSumProblem {
405        expr: String,
406        a: Tensor,
407        b: Tensor,
408    }
409
410    impl EinSumProblem {
411        fn check(&self) -> TractResult<()> {
412            let mut model = TypedModel::default();
413            let sa = model.add_source("a", f32::fact(self.a.shape()))?;
414            let sb = model.add_source("b", f32::fact(self.b.shape()))?;
415            let einsum = model.wire_node(
416                "einsum",
417                EinSum::new(self.expr.parse().unwrap(), f32::datum_type()),
418                &[sa, sb],
419            )?;
420            model.select_output_outlets(&einsum)?;
421            let a = self.a.clone().into_tvalue();
422            let b = self.b.clone().into_tvalue();
423            let inputs = tvec!(a, b);
424            let reference = TypedRunnableModel::new(model.clone())?.run(inputs.clone())?.remove(0);
425            rewrite_einsum_to_prefix_matmul(&mut model, true)?;
426            assert!(model.nodes.iter().all(|n| !n.op_is::<EinSum>()));
427            let test = TypedRunnableModel::new(model)?.run(inputs)?.remove(0);
428            reference.close_enough(&test, true)
429        }
430    }
431
432    #[rustfmt::skip] #[test] fn prop_mk_kn_mn() -> TestCaseResult { test_expr("mk,kn->mn") }
433    #[rustfmt::skip] #[test] fn prop_km_kn_mn() -> TestCaseResult { test_expr("km,kn->mn") }
434    #[rustfmt::skip] #[test] fn prop_mk_nk_mn() -> TestCaseResult { test_expr("mk,nk->mn") }
435    #[rustfmt::skip] #[test] fn prop_mk_kn_nm() -> TestCaseResult { test_expr("mk,kn->nm") }
436    #[rustfmt::skip] #[test] fn prop_k_kn_mn() -> TestCaseResult { test_expr("k,kn->mn") }
437    #[rustfmt::skip] #[test] fn prop_mk_k_mn() -> TestCaseResult { test_expr("mk,k->mn") }
438    #[rustfmt::skip] #[test] fn prop_m_n_mn() -> TestCaseResult { test_expr("m,n->mn") }
439    #[rustfmt::skip] #[test] fn prop_amk_akn_amn() -> TestCaseResult { test_expr("amk,akn->amn") }
440    #[rustfmt::skip] #[test] fn prop_mk_akn_amn() -> TestCaseResult { test_expr("mk,akn->amn") }
441    #[rustfmt::skip] #[test] fn prop_btgi_gih_tgh() -> TestCaseResult { test_expr("btgi,gih->tgh") }
442    #[rustfmt::skip] #[test] fn prop_tgi_gih_btgh() -> TestCaseResult { test_expr("tgi,gih->btgh") }
443
444    #[test]
445    fn k_kn_mn_0() -> TractResult<()> {
446        EinSumProblem {
447            expr: "k,kn->mn".to_string(),
448            a: tensor1(&[0f32, 0f32]),
449            b: tensor2(&[[0f32, 0.], [0., 0.]]),
450        }
451        .check()
452    }
453
454    #[test]
455    fn mk_k_mn_0() -> TractResult<()> {
456        EinSumProblem {
457            expr: "mk,k->mn".to_string(),
458            a: Tensor::zero::<f32>(&[2, 2]).unwrap(),
459            b: Tensor::zero::<f32>(&[2]).unwrap(),
460        }
461        .check()
462    }
463
464    #[test]
465    fn mk_k_mn_1() -> TractResult<()> {
466        EinSumProblem {
467            expr: "mk,k->mn".to_string(),
468            a: Tensor::zero::<f32>(&[1, 2]).unwrap(),
469            b: Tensor::zero::<f32>(&[2]).unwrap(),
470        }
471        .check()
472    }
473
474    #[test]
475    fn mk_kn_nm_0() -> TractResult<()> {
476        EinSumProblem {
477            expr: "mk,kn->mn".to_string(),
478            a: Tensor::zero::<f32>(&[3, 2]).unwrap(),
479            b: Tensor::zero::<f32>(&[2, 2]).unwrap(),
480        }
481        .check()
482    }
483
484    #[test]
485    fn amk_akn_amn_0() -> TractResult<()> {
486        EinSumProblem {
487            expr: "amk,akn->amn".to_string(),
488            a: Tensor::zero::<f32>(&[1, 1, 2]).unwrap(),
489            b: Tensor::zero::<f32>(&[1, 2, 1]).unwrap(),
490        }
491        .check()
492    }
493
494    #[test]
495    fn amk_akn_amn_1() -> TractResult<()> {
496        EinSumProblem {
497            expr: "amk,akn->amn".to_string(),
498            a: Tensor::zero::<f32>(&[2, 1, 2]).unwrap(),
499            b: Tensor::zero::<f32>(&[1, 2, 1]).unwrap(),
500        }
501        .check()
502    }
503
504    #[test]
505    fn amk_akn_amn_2() -> TractResult<()> {
506        EinSumProblem {
507            expr: "amk,akn->amn".to_string(),
508            a: Tensor::zero::<f32>(&[1, 1, 2]).unwrap(),
509            b: Tensor::zero::<f32>(&[2, 2, 2]).unwrap(),
510        }
511        .check()
512    }
513
514    #[test]
515    fn amk_akn_amn_3() -> TractResult<()> {
516        EinSumProblem {
517            expr: "amk,akn->amn".to_string(),
518            a: Tensor::zero::<f32>(&[1, 1, 2]).unwrap(),
519            b: Tensor::zero::<f32>(&[2, 2, 1]).unwrap(),
520        }
521        .check()
522    }
523
524    #[test]
525    fn km_anbck_bmn_0() -> TractResult<()> {
526        EinSumProblem {
527            expr: "km,anbck->bmn".to_string(),
528            a: Tensor::zero::<f32>(&[2, 1]).unwrap(),
529            b: Tensor::zero::<f32>(&[1, 1, 1, 1, 2]).unwrap(),
530        }
531        .check()
532    }
533
534    fn check_k1(expr: &str, a_shape: &[usize], b_shape: &[usize]) -> TractResult<()> {
535        let a_len = a_shape.iter().product::<usize>();
536        let b_len = b_shape.iter().product::<usize>();
537        let a =
538            tensor1(&(0..a_len).map(|i| (i + 1) as f32).collect_vec()).into_shape(a_shape).unwrap();
539        let b = tensor1(&(0..b_len).map(|i| (i + 7) as f32 * 0.5).collect_vec())
540            .into_shape(b_shape)
541            .unwrap();
542        EinSumProblem { expr: expr.to_string(), a, b }.check()
543    }
544
545    // K=1 means the einsum's contraction degenerates to broadcast-mul. detect_rule
546    // short-circuits to a Mul op rather than dispatching the GEMM kernel for a single
547    // FMA per tile. The gmk,Ngnk->Ngmn pattern with k=1 is what depthwise ConvTranspose
548    // lowers to, and the per-tile GEMM overhead dominates in that case.
549    #[test]
550    fn k1_amk_akn_amn() -> TractResult<()> {
551        check_k1("amk,akn->amn", &[2, 3, 1], &[2, 1, 4])
552    }
553
554    #[test]
555    fn k1_gmk_ngnk_ngmn() -> TractResult<()> {
556        check_k1("gmk,Ngnk->Ngmn", &[3, 2, 1], &[1, 3, 4, 1])
557    }
558
559    #[test]
560    fn k1_mk_kn_mn() -> TractResult<()> {
561        check_k1("mk,kn->mn", &[2, 1], &[1, 3])
562    }
563
564    // The unit_k_to_broadcast_mul declutter rule is scoped to einsums whose output is
565    // consumed by DeconvSum (the original ConvTranspose-K=1 target). For these isolated
566    // einsum tests the rule won't fire — outputs still go through OptMatMul correctly,
567    // just slower than a broadcast Mul would be. The end-to-end ConvTranspose case is
568    // covered by integration tests (DFN3 erb_dec, GTCRN) and verified bit-exact in the
569    // PR description.
570    #[test]
571    fn k1_no_k_axis_outer() -> TractResult<()> {
572        check_k1("m,n->mn", &[3], &[4])
573    }
574
575    #[test]
576    fn k1_high_rank_no_decov_sum() -> TractResult<()> {
577        // From a CI proptest failure: a high-rank einsum with K=1, no DeconvSum
578        // downstream → rule must not fire, OptMatMul handles it correctly.
579        check_k1("wmexk,wxnk->ewnxm", &[2, 1, 2, 2, 1], &[2, 2, 1, 1])
580    }
581
582    #[test]
583    fn k1_sdpa_shape_no_deconv_sum() -> TractResult<()> {
584        // Regression for SDPA failure mode: when head_dim=1, SDPA's score einsum
585        // looks like a K=1 case structurally, but it's followed by Softmax (not
586        // DeconvSum). Rule must NOT fire — Metal's SDPA fusion would otherwise break.
587        check_k1("bhmk,bhnk->bhmn", &[1, 3, 4, 1], &[1, 3, 4, 1])
588    }
589
590    #[test]
591    fn q() -> TractResult<()> {
592        let qp = QParams::ZpScale { zero_point: 0, scale: 0.1 };
593        let op = EinSum {
594            axes: "mk,kn,m,,,,,,->mn".parse()?,
595            operating_dt: i32::datum_type(),
596            q_params: Some(DatumType::QI8(qp)),
597        };
598        let mut model = TypedModelPatch::default();
599        let inputs = [
600            model.add_source("a", DatumType::QI8(qp).fact([3, 2]))?,
601            model.add_source("b", DatumType::QI8(qp).fact([2, 4]))?,
602            model.add_source("bias", i32::datum_type().fact([3]))?,
603            model.add_const("a0", tensor0(qp.zp_scale().0))?,
604            model.add_const("a_scale", tensor0(qp.zp_scale().1))?,
605            model.add_const("b0", tensor0(qp.zp_scale().0))?,
606            model.add_const("b_scale", tensor0(qp.zp_scale().1))?,
607            model.add_const("c0", tensor0(qp.zp_scale().0))?,
608            model.add_const("c_scale", tensor0(qp.zp_scale().1))?,
609        ];
610        let wire = model.wire_node("einsum", op.clone(), &inputs)?;
611        model.select_output_outlets(&wire)?;
612        rewrite_einsum_to_prefix_matmul(&mut model, true)?;
613        assert!(model.nodes.iter().all(|n| !n.op_is::<EinSum>()));
614        Ok(())
615    }
616}
tract_core/ops/einsum/prefix_matmul.rs

tract_core/ops/einsum/
prefix_matmul.rs
