tract_core/ops/einsum/
prefix_matmul.rs

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

tract_core/ops/einsum/
prefix_matmul.rs
