Skip to main content

tract_core/ops/
binary.rs

1use crate::internal::*;
2use downcast_rs::Downcast;
3use dyn_eq::DynEq;
4use std::fmt::{self, Debug};
5use tract_data::itertools::izip;
6use tract_itertools::Itertools;
7use tract_linalg::multithread::BShare;
8use tract_linalg::{BinOp, LinalgFn};
9
10use super::math::{Add, Max, Min, Mul, Sub};
11use super::{cast::cast, math::SubF};
12
13pub trait BinMiniOp:
14    fmt::Debug + dyn_clone::DynClone + dyn_eq::DynEq + Send + Sync + 'static + Downcast
15{
16    fn name(&self) -> &'static str;
17    fn validation(&self) -> Validation {
18        Validation::Accurate
19    }
20    fn operating_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType> {
21        a.common_super_type(b).with_context(|| format_err!("No super type for {:?} and {:?}", a, b))
22    }
23    fn result_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType>;
24    fn eval_in_a(&self, a: &mut Tensor, b: &Tensor) -> TractResult<()>;
25    fn eval_out_of_place(&self, c: &mut Tensor, a: &Tensor, b: &Tensor) -> TractResult<()>;
26
27    fn is_commutative(&self) -> bool {
28        true
29    }
30    fn neutral_element(&self) -> Option<i64> {
31        None
32    }
33    fn absorbing_element(&self) -> Option<i64> {
34        None
35    }
36
37    #[allow(unused_variables)]
38    fn maybe_eval_qbinary_as_float_op(
39        &self,
40        a: &TValue,
41        b: &TValue,
42        c_dt: &DatumType,
43    ) -> TractResult<Option<Tensor>> {
44        Ok(None)
45    }
46
47    fn generic_eval(&self, a: TValue, b: TValue, c_dt: DatumType) -> TractResult<Tensor> {
48        if let Some(tensor) = self.maybe_eval_qbinary_as_float_op(&a, &b, &c_dt)? {
49            return Ok(tensor);
50        }
51        // Same-shape fast path: skip `multi_broadcast` allocation when shapes
52        // are already equal (very common: residuals, mask application, etc.).
53        // Correctness: equal shapes imply broadcast shape == a.shape() and the
54        // existing slow path would have taken this same branch.
55        if c_dt == a.datum_type() && a.shape() == b.shape() {
56            let mut a = a.into_tensor();
57            self.eval_in_a(&mut a, &b)?;
58            return Ok(a);
59        }
60        let c_shape = crate::broadcast::multi_broadcast(&[a.shape(), b.shape()])?;
61        if &*c_shape == a.shape() && c_dt == a.datum_type() {
62            let mut a = a.into_tensor();
63            self.eval_in_a(&mut a, &b)?;
64            Ok(a)
65        } else {
66            let mut c = unsafe { Tensor::uninitialized_dt(c_dt, &c_shape)? };
67            self.eval_out_of_place(&mut c, &a, &b)?;
68            Ok(c)
69        }
70    }
71    fn eval(&self, a: TValue, b: TValue, c_dt: DatumType) -> TractResult<Tensor> {
72        self.generic_eval(a, b, c_dt)
73    }
74    #[allow(unused_variables)]
75    fn declutter(
76        &self,
77        model: &TypedModel,
78        node: &TypedNode,
79    ) -> TractResult<Option<TypedModelPatch>> {
80        Ok(None)
81    }
82    #[allow(unused_variables)]
83    fn codegen(
84        &self,
85        model: &TypedModel,
86        node: &TypedNode,
87    ) -> TractResult<Option<TypedModelPatch>> {
88        Ok(None)
89    }
90    #[allow(unused_variables)]
91    fn cost_per_element(&self, dt: DatumType) -> TVec<(Cost, usize)> {
92        tvec!()
93    }
94    fn as_linalg_binop(&self) -> Option<tract_linalg::BinOp> {
95        None
96    }
97
98    /// Override for ops that can evaluate symbolic TDim inputs (comparisons).
99    #[allow(unused_variables)]
100    fn eval_symbolic(
101        &self,
102        session: &TurnState,
103        inputs: TVec<TValue>,
104    ) -> TractResult<Option<TVec<TValue>>> {
105        Ok(None)
106    }
107
108    /// Override for ops that produce TDim-level comparison expressions (comparisons).
109    #[allow(unused_variables)]
110    fn uniform_tdim_comparison(&self, a: &TDim, b: &TDim) -> Option<TDim> {
111        None
112    }
113}
114dyn_clone::clone_trait_object!(BinMiniOp);
115dyn_eq::eq_trait_object!(BinMiniOp);
116downcast_rs::impl_downcast!(BinMiniOp);
117
118#[derive(Debug, Clone, PartialEq, Eq)]
119pub struct TypedBinOp(pub Box<dyn BinMiniOp>, pub Option<DatumType>);
120
121impl Op for TypedBinOp {
122    fn name(&self) -> StaticName {
123        self.0.name().into()
124    }
125
126    fn validation(&self) -> Validation {
127        self.0.validation()
128    }
129
130    op_as_typed_op!();
131}
132
133impl TypedBinOp {
134    fn output_datum_type(&self, a_dt: DatumType, b_dt: DatumType) -> TractResult<DatumType> {
135        if let Some(dt) = self.1 { Ok(dt) } else { self.0.result_datum_type(a_dt, b_dt) }
136    }
137}
138
139impl EvalOp for TypedBinOp {
140    fn is_stateless(&self) -> bool {
141        true
142    }
143
144    fn eval_with_session(
145        &self,
146        _node_id: usize,
147        session: &TurnState,
148        inputs: TVec<TValue>,
149    ) -> TractResult<TVec<TValue>> {
150        if let Some(result) = self.0.eval_symbolic(session, inputs.clone())? {
151            return Ok(result);
152        }
153        let (a, b) = args_2!(inputs);
154        ensure!(a.rank() == b.rank());
155        let c_dt = self.output_datum_type(a.datum_type(), b.datum_type())?;
156        Ok(tvec!(self.0.eval(a, b, c_dt)?.into_tvalue()))
157    }
158
159    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
160        let (a, b) = args_2!(inputs);
161        ensure!(a.rank() == b.rank());
162        let c_dt = self.output_datum_type(a.datum_type(), b.datum_type())?;
163        Ok(tvec!(self.0.eval(a, b, c_dt)?.into_tvalue()))
164    }
165}
166
167impl TypedBinOp {
168    fn combine_uniform_tdim(&self, a: &TDim, b: &TDim) -> Option<TDim> {
169        // Comparison ops provide their own TDim combination
170        if let Some(result) = self.0.uniform_tdim_comparison(a, b) {
171            return Some(result);
172        }
173        let a = tensor0(a.clone()).into_tvalue();
174        let b = tensor0(b.clone()).into_tvalue();
175        let result = self.0.eval(a, b, TDim::datum_type()).ok()?;
176        result
177            .try_as_plain()
178            .ok()
179            .and_then(|d| d.as_slice::<TDim>().ok())
180            .and_then(|s| s.first())
181            .cloned()
182            .map(|d| d.reduce())
183    }
184
185    fn combine_uniform_tdim_with_konst(&self, a: &TDim, konst: &Tensor) -> Option<TDim> {
186        if konst.len() != 1 {
187            return None;
188        }
189        // Integer-valued scalar (including float constants like 2.0, 1.0, 3.0)
190        let b_int: Option<i64> =
191            if konst.datum_type().is_integer() || konst.datum_type().is::<bool>() {
192                konst.cast_to_scalar::<i64>().ok()
193            } else if konst.datum_type().is_float() {
194                konst.cast_to_scalar::<f64>().ok().and_then(|f| {
195                    if (f - f.round()).abs() < 1e-6 { Some(f.round() as i64) } else { None }
196                })
197            } else {
198                None
199            };
200        if let Some(b) = b_int {
201            return self.combine_uniform_tdim(a, &TDim::Val(b));
202        }
203        // Mul by reciprocal of integer (e.g. ×0.5 → Div(a, 2))
204        if self.0.neutral_element() == Some(1)
205            && let Some(f) = konst.cast_to_scalar::<f64>().ok().filter(|&f| f > 0.0)
206        {
207            let n = (1.0 / f).round() as u64;
208            if n >= 2 && (f * n as f64 - 1.0).abs() < 1e-6 {
209                return Some(TDim::Div(Box::new(a.clone()), n).reduce());
210            }
211        }
212        None
213    }
214}
215
216impl TypedOp for TypedBinOp {
217    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
218        if inputs[0].rank() != inputs[1].rank() {
219            bail!(
220                "Typed ops require rank match. Invalid inputs for {}: {}",
221                self.name(),
222                inputs.iter().map(|s| format!("{s:?}")).join(" ; ")
223            );
224        }
225        let out_dt = self.output_datum_type(inputs[0].datum_type, inputs[1].datum_type)?;
226        let mut fact = out_dt.fact(&*crate::broadcast::multi_broadcast(&[
227            &inputs[0].shape.to_tvec(),
228            &inputs[1].shape.to_tvec(),
229        ])?);
230        if let (Some(a), Some(b)) = (&inputs[0].uniform_tdim, &inputs[1].uniform_tdim) {
231            fact.uniform_tdim = self.combine_uniform_tdim(a, b);
232            // And(a,b) has no TDim kernel; for 0/1 booleans And == Mul
233            if fact.uniform_tdim.is_none() && self.0.is::<crate::ops::logic::And>() {
234                fact.uniform_tdim = Some(TDim::Mul(vec![a.clone(), b.clone()]).reduce());
235            }
236        }
237        // Fallback: one side has uniform_tdim, the other is a scalar constant
238        if fact.uniform_tdim.is_none() {
239            for (expr, konst_fact) in [
240                (inputs[0].uniform_tdim.as_ref(), inputs[1]),
241                (inputs[1].uniform_tdim.as_ref(), inputs[0]),
242            ] {
243                let Some(a) = expr else { continue };
244                let Some(konst) = konst_fact.konst.as_ref() else { continue };
245                fact.uniform_tdim = self.combine_uniform_tdim_with_konst(a, konst);
246                if fact.uniform_tdim.is_some() {
247                    break;
248                }
249            }
250        }
251        Ok(tvec!(fact))
252    }
253
254    fn input_roi(
255        &self,
256        model: &TypedModel,
257        node: &TypedNode,
258    ) -> TractResult<Option<TVec<Option<TDim>>>> {
259        // Introduction: Mul (or any op with neutral_element=1) with a mask
260        // that has uniform_tdim → the other input gets that expression as ROI.
261        if self.0.neutral_element() == Some(1) {
262            for (mask_ix, other_ix) in [(0usize, 1usize), (1, 0)] {
263                let fact = model.outlet_fact(node.inputs[mask_ix])?;
264                if let Some(mask_expr) = &fact.uniform_tdim {
265                    let mut rois = tvec![None; node.inputs.len()];
266                    rois[other_ix] = Some(mask_expr.clone());
267                    return Ok(Some(rois));
268                }
269            }
270        }
271        // Bubbling: delegate to the natural blanket implementation.
272        crate::optim::propagate_roi::bubble_roi(model, node)
273    }
274
275    fn change_axes(
276        &self,
277        model: &TypedModel,
278        node: &TypedNode,
279        _io: InOut,
280        change: &AxisOp,
281    ) -> TractResult<Option<AxisChangeConsequence>> {
282        if let AxisOp::Rm(rm) = change {
283            let (inputs, outputs) = model.node_facts(node.id)?;
284            if inputs.len() >= 2
285                && outputs.len() >= 1
286                && inputs[0].rank() > *rm
287                && inputs[1].rank() > *rm
288                && outputs[0].rank() > *rm
289            {
290                rule_if!(inputs[0].shape[*rm].is_one());
291                rule_if!(inputs[1].shape[*rm].is_one());
292                rule_if!(outputs[0].shape[*rm].is_one());
293            }
294        }
295        Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
296    }
297
298    fn axes_mapping(
299        &self,
300        inputs: &[&TypedFact],
301        outputs: &[&TypedFact],
302    ) -> TractResult<AxesMapping> {
303        AxesMapping::natural(inputs, outputs)
304    }
305
306    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
307        let count: TDim = self.output_facts(inputs)?[0].shape.iter().product();
308        Ok(self
309            .0
310            .cost_per_element(inputs[0].datum_type)
311            .into_iter()
312            .map(|(c, n)| (c, count.clone() * n))
313            .collect())
314    }
315
316    fn slice(
317        &self,
318        patch: &mut TypedModelPatch,
319        _model: &TypedModel,
320        _node: &TypedNode,
321        prefix: &str,
322        inputs: &[OutletId],
323        _output_axis: usize,
324        _start: &TDim,
325        _end: &TDim,
326    ) -> TractResult<Option<TVec<OutletId>>> {
327        Ok(Some(patch.wire_node(prefix, self.clone(), inputs)?))
328    }
329
330    fn declutter(
331        &self,
332        model: &TypedModel,
333        node: &TypedNode,
334    ) -> TractResult<Option<TypedModelPatch>> {
335        let (a_dt, b_dt) = if let &[a, b] = &*model.node_input_facts(node.id)? {
336            (a.datum_type().unwrap(), b.datum_type().unwrap())
337        } else {
338            unreachable!("TypedBinOp has two inputs.")
339        };
340        if let Some(neutral_patch) =
341            declutter_neutral(model, node, self.0.as_ref(), self.output_datum_type(a_dt, b_dt)?)?
342        {
343            return Ok(Some(neutral_patch));
344        }
345        if let Some(absorbing_patch) = declutter_absorbing(model, node, self.0.as_ref())? {
346            return Ok(Some(absorbing_patch));
347        }
348        if let Some(broadcast_patch) =
349            declutter_broadcasting_operand_1(model, node, self.0.clone())?
350        {
351            return Ok(Some(broadcast_patch));
352        }
353        self.0.declutter(model, node)
354    }
355
356    fn codegen(
357        &self,
358        model: &TypedModel,
359        node: &TypedNode,
360    ) -> TractResult<Option<TypedModelPatch>> {
361        if let Some(linalg_bin_op) = self.0.as_linalg_binop() {
362            let input_facts = model.node_input_facts(node.id)?;
363            let must_swap_inputs =
364                input_facts.iter().collect_tuple().is_some_and(|(a_fact, b_fact)| {
365                    (a_fact.shape.volume() - b_fact.shape.volume()).prove_strict_negative()
366                });
367            let (operand_1, operand_2) = if must_swap_inputs {
368                (input_facts[1], input_facts[0])
369            } else {
370                (input_facts[0], input_facts[1])
371            };
372
373            let (by_scalar_should_be_efficient, unicast_should_be_efficient) =
374                find_most_efficient_config(model, node, must_swap_inputs)?;
375
376            // Check if op is quantized
377            let c_dt = self.output_datum_type(operand_1.datum_type, operand_2.datum_type)?;
378            let op_is_quant = c_dt.is_quantized()
379                || operand_1.datum_type.is_quantized()
380                || operand_2.datum_type.is_quantized();
381
382            // Check if it can be evaluated in a
383            let c_dt = self.output_datum_type(operand_1.datum_type, operand_2.datum_type)?;
384            let c_shape = crate::broadcast::multi_broadcast(&[
385                operand_1.shape.clone(),
386                operand_2.shape.clone(),
387            ])?;
388            let can_eval_in_a =
389                (c_shape.to_vec() == operand_1.shape.to_vec()) && (c_dt == operand_1.datum_type);
390
391            // Swap input if required
392            let inputs = if must_swap_inputs {
393                let mut swap_input = node.inputs.clone();
394                swap_input.swap(0, 1);
395                swap_input
396            } else {
397                node.inputs.clone()
398            };
399            let actual_linalg_op =
400                if must_swap_inputs { linalg_bin_op.flip() } else { linalg_bin_op };
401            let actual_core_op = core_op_for_linalg_op(&actual_linalg_op);
402
403            let dt = model.node_input_facts(node.id)?[0].datum_type;
404            if by_scalar_should_be_efficient & can_eval_in_a & !op_is_quant {
405                rule_if_some!(func = tract_linalg::bin_by_scalar(dt, actual_linalg_op));
406                let eval_fn = Arc::from(func);
407                return Ok(Some(
408                    TypedModelPatch::replace_single_op(
409                        model,
410                        node,
411                        &inputs,
412                        OptBinByScalar { binop: actual_core_op, eval_fn },
413                    )?
414                    .with_context("ByScalar"),
415                ));
416            }
417
418            if unicast_should_be_efficient & can_eval_in_a & !op_is_quant {
419                rule_if_some!(func = tract_linalg::bin_unicast(dt, actual_linalg_op));
420                let eval_fn = Arc::from(func);
421                return Ok(Some(
422                    TypedModelPatch::replace_single_op(
423                        model,
424                        node,
425                        &inputs,
426                        OptBinUnicast { binop: actual_core_op, eval_fn },
427                    )?
428                    .with_context("Unicast"),
429                ));
430            }
431        }
432
433        Ok(None)
434    }
435    as_op!();
436}
437
438fn core_op_for_linalg_op(linalg: &BinOp) -> Box<dyn BinMiniOp> {
439    match linalg {
440        BinOp::Min => Box::new(Min),
441        BinOp::Max => Box::new(Max),
442        BinOp::Add => Box::new(Add),
443        BinOp::Mul => Box::new(Mul),
444        BinOp::Sub => Box::new(Sub),
445        BinOp::SubF => Box::new(SubF),
446    }
447}
448fn declutter_broadcasting_operand_1(
449    model: &TypedModel,
450    node: &TypedNode,
451    mini_op: Box<dyn BinMiniOp>,
452) -> TractResult<Option<TypedModelPatch>> {
453    let (a_shape, b_shape) = if let &[a, b] = &*model.node_input_facts(node.id)? {
454        (a.shape.clone(), b.shape.clone())
455    } else {
456        unreachable!("TypedBinOp has two inputs.")
457    };
458
459    let a_num_elements = a_shape.iter().product::<TDim>();
460    let b_num_elements = b_shape.iter().product::<TDim>();
461    let a_should_be_broadcast = (a_num_elements - b_num_elements).prove_strict_negative();
462    if a_should_be_broadcast & mini_op.is_commutative() {
463        let mut swap_input = node.inputs.clone();
464        swap_input.swap(0, 1);
465        return Ok(Some(TypedModelPatch::replace_single_op(
466            model,
467            node,
468            &swap_input,
469            TypedBinOp(mini_op, None),
470        )?));
471    }
472
473    Ok(None)
474}
475
476/// Shunt a binary op whose uniform input holds the op's neutral element
477/// (`x + 0`, `x * 1`, `x - 0`).
478///
479/// The neutral element is compared on dequantized values, so a quantized node
480/// can be arithmetically neutral while still re-encoding its input: `out_dt`
481/// may carry other quantization parameters than the variable input, and the
482/// same real value is then a different integer. Such a node degrades to a
483/// `Cast` instead of vanishing.
484fn declutter_neutral(
485    model: &TypedModel,
486    node: &TypedNode,
487    mini_op: &dyn BinMiniOp,
488    out_dt: DatumType,
489) -> TractResult<Option<TypedModelPatch>> {
490    let Some(uniform) = crate::ops::binary::one_input_is_uniform(model, node)? else {
491        return Ok(None);
492    };
493    let uni_is_neutral = mini_op
494        .neutral_element()
495        .is_some_and(|neutral| tensor0(neutral).close_enough(&uniform.uni, false).is_ok());
496    // Non-commutative ops only have a right neutral: x - 0 == x, but 0 - x == -x.
497    let uniform_on_neutral_side = mini_op.is_commutative() || !uniform.left_is_uniform;
498    if !uni_is_neutral || !uniform_on_neutral_side {
499        return Ok(None);
500    }
501    if uniform.uni.datum_type().is_quantized() {
502        return Ok(Some(TypedModelPatch::replace_single_op(
503            model,
504            node,
505            &[uniform.var],
506            cast(out_dt),
507        )?));
508    }
509    Ok(Some(TypedModelPatch::rewire(model, &[uniform.var], &[node.id.into()], &|_, inputs| {
510        Ok(inputs.into())
511    })?))
512}
513
514/// When one input is the absorbing element (e.g. 0 for Mul, false for And),
515/// replace the entire op with a uniform-value tensor of the output shape.
516///
517/// We can't shunt the uniform input directly: it may be lower-rank or have
518/// broadcast-from-1 dims that don't match the op's output shape (e.g.
519/// `Mul([4, 1], scalar-0)` outputs `[4, 1]`, not `[1]`).  Wire a
520/// `MultiBroadcastTo` from the uniform constant to the output shape;
521/// subsequent declutter folds it into a pure constant when the shape is
522/// fully concrete.
523fn declutter_absorbing(
524    model: &TypedModel,
525    node: &TypedNode,
526    mini_op: &dyn BinMiniOp,
527) -> TractResult<Option<TypedModelPatch>> {
528    if let Some(uniform) = crate::ops::binary::one_input_is_uniform(model, node)? {
529        let is_absorbing = mini_op
530            .absorbing_element()
531            .map(|absorb| tensor0(absorb).close_enough(&uniform.uni, false).is_ok())
532            .unwrap_or(false);
533        if is_absorbing {
534            let output_fact = model.outlet_fact(node.id.into())?;
535            let output_dt = output_fact.datum_type;
536            let output_shape = output_fact.shape.clone();
537            let uni_inlet = if uniform.left_is_uniform { 0 } else { 1 };
538            let uni_input_shape = &model.outlet_fact(node.inputs[uni_inlet])?.shape;
539            // Fast path: shapes and types match — shunt the absorbing input directly.
540            if uni_input_shape == &output_shape && uniform.uni.datum_type() == output_dt {
541                return Ok(Some(TypedModelPatch::rewire(
542                    model,
543                    &[node.inputs[uni_inlet]],
544                    &[node.id.into()],
545                    &|_, inputs| Ok(inputs.into()),
546                )?));
547            }
548            // General path: create a constant encoded in the output type.
549            // This handles both shape mismatches and quantization mismatches
550            // (e.g. absorbing input is QU8(Z:61 S:1) but output is QU8(Z:0 S:0.5)).
551            let absorb_val = mini_op.absorbing_element().unwrap();
552            let absorbing_const =
553                tensor0(absorb_val as f32).cast_to_dt(output_dt)?.into_owned().into_arc_tensor();
554            let mut patch = TypedModelPatch::default();
555            let uni_const =
556                patch.add_const(format!("{}.absorbing_const", node.name), absorbing_const)?;
557            let bcast = patch.wire_node(
558                format!("{}.absorbing_bcast", node.name),
559                crate::ops::array::MultiBroadcastTo { shape: output_shape },
560                &[uni_const],
561            )?[0];
562            patch.shunt_outside(model, node.id.into(), bcast)?;
563            return Ok(Some(patch));
564        }
565    }
566    Ok(None)
567}
568
569fn find_most_efficient_config(
570    model: &TypedModel,
571    node: &TypedNode,
572    swap_input: bool,
573) -> TractResult<(bool, bool)> {
574    if let &[a, b] = &*model.node_input_facts(node.id)? {
575        let a_shape = if swap_input { b.shape.clone() } else { a.shape.clone() };
576        let b_shape = if swap_input { a.shape.clone() } else { b.shape.clone() };
577
578        let by_scalar_is_possible = OptBinByScalar::check_input_shapes(&a_shape, &b_shape);
579        let num_by_scalar_elements = if by_scalar_is_possible {
580            a_shape
581                .iter()
582                .zip(b_shape.iter())
583                .rev()
584                .take_while(|(_, rev_b_dim)| **rev_b_dim == TDim::Val(1))
585                .map(|(rev_a_dim, _)| rev_a_dim)
586                .product::<TDim>()
587        } else {
588            TDim::Val(0)
589        };
590
591        let unicast_is_possible = OptBinUnicast::check_input_shapes(&a_shape, &b_shape);
592        let num_unicast_elements = if unicast_is_possible {
593            a_shape
594                .iter()
595                .zip(b_shape.iter())
596                .rev()
597                .take_while(|(a_dim, b_dim)| a_dim == b_dim)
598                .map(|(a_dim, _)| a_dim)
599                .product::<TDim>()
600        } else {
601            TDim::Val(0)
602        };
603
604        let min_num_elements = 32;
605        let by_scalar_should_be_efficient = gt_tdim(num_by_scalar_elements, min_num_elements);
606        let unicast_should_be_efficient = gt_tdim(num_unicast_elements, min_num_elements);
607        return Ok((by_scalar_should_be_efficient, unicast_should_be_efficient));
608    }
609    Ok((false, false))
610}
611
612pub fn gt_tdim(x: TDim, min_val: i64) -> bool {
613    TDim::Val(min_val).mini(x).as_i64().is_some_and(|v| v == min_val)
614}
615
616#[derive(Clone)]
617pub struct OptBinByScalar {
618    pub binop: Box<dyn BinMiniOp>,
619    eval_fn: Arc<LinalgFn>,
620}
621
622impl Debug for OptBinByScalar {
623    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
624        f.debug_struct("OptBinByScalar").field("binop", &self.binop).finish()
625    }
626}
627
628impl OptBinByScalar {
629    fn check_input_shapes(a_shape: &[TDim], b_shape: &[TDim]) -> bool {
630        if a_shape.len() != b_shape.len() {
631            return false;
632        };
633
634        a_shape
635            .iter()
636            .zip(b_shape.iter())
637            .skip_while(|(a_dim, b_dim)| a_dim == b_dim)
638            .all(|(_, b_dim)| *b_dim == 1.to_dim())
639    }
640}
641
642impl PartialEq for OptBinByScalar {
643    fn eq(&self, other: &Self) -> bool {
644        *self.binop == *other.binop
645    }
646}
647impl Eq for OptBinByScalar {}
648
649impl Op for OptBinByScalar {
650    fn name(&self) -> StaticName {
651        format!("Opt{}ByScalar", self.binop.name()).into()
652    }
653
654    op_as_typed_op!();
655}
656
657impl EvalOp for OptBinByScalar {
658    fn is_stateless(&self) -> bool {
659        true
660    }
661
662    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
663        let (a, b) = args_2!(inputs);
664        // Same as OptBinUnicast: the fast path uses at_prefix + as_slice_mut
665        // and relies on natural C-order strides for the slice math. Fall back
666        // to the generic eval if either operand has non-natural strides or a
667        // storage size that doesn't match its declared shape (e.g. after
668        // Tensor::insert_axis which leaves non-natural strides behind).
669        let a_natural = a.len() == a.shape().iter().product::<usize>()
670            && a.strides() == &*Tensor::natural_strides(a.shape());
671        let b_natural = b.len() == b.shape().iter().product::<usize>()
672            && b.strides() == &*Tensor::natural_strides(b.shape());
673        if !a_natural || !b_natural {
674            let c_dt = self.binop.result_datum_type(a.datum_type(), b.datum_type())?;
675            return Ok(tvec!(self.binop.eval(a, b, c_dt)?.into_tvalue()));
676        }
677
678        let mut a = a.into_tensor();
679        let b_shape = b.shape();
680
681        let first_unary_axis = b_shape
682            .iter()
683            .enumerate()
684            .rev()
685            .take_while(|&(_, &dim)| dim == 1)
686            .map(|(i, _)| i)
687            .last()
688            .context("Cannot use by_scalar when no trailing dimensions are unary")?;
689
690        // b carries one scalar per block of a, the blocks being a's axes before
691        // first_unary_axis; check_input_shapes makes b's match a's there.
692        let n_blocks: usize = a.shape()[..first_unary_axis].iter().product();
693        // A zero-sized dim zeroes n_blocks, and par_bin no-ops on an empty a.
694        let period = a.len().checked_div(n_blocks).unwrap_or(0);
695        tract_linalg::multithread::par_bin(&*self.eval_fn, &mut a, &b, period, BShare::PerBlock)?;
696        Ok(tvec!(a.into_tvalue()))
697    }
698}
699
700impl TypedOp for OptBinByScalar {
701    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
702        ensure!(Self::check_input_shapes(&inputs[0].shape, &inputs[1].shape));
703        let out_dt = self.binop.result_datum_type(inputs[0].datum_type, inputs[1].datum_type)?;
704        let out_shape = inputs[0].shape.clone();
705        Ok(tvec!(out_dt.fact(out_shape)))
706    }
707
708    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
709        let count: TDim = self.output_facts(inputs)?[0].shape.iter().product();
710        Ok(self
711            .binop
712            .cost_per_element(inputs[0].datum_type)
713            .into_iter()
714            .map(|(c, n)| (c, count.clone() * n))
715            .collect())
716    }
717
718    as_op!();
719}
720
721#[derive(Clone)]
722pub struct OptBinUnicast {
723    pub binop: Box<dyn BinMiniOp>,
724    eval_fn: Arc<LinalgFn>,
725}
726
727impl Debug for OptBinUnicast {
728    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
729        f.debug_struct("OptBinUnicast").field("binop", &self.binop).finish()
730    }
731}
732
733impl OptBinUnicast {
734    fn check_b_alignement(a_shape: &[TDim], b_shape: &[TDim]) -> bool {
735        let num_iterations: TDim = a_shape
736            .iter()
737            .zip(b_shape.iter())
738            .take_while(|(_, b_dim)| **b_dim == 1.to_dim())
739            .map(|(a_dim, _)| a_dim)
740            .product();
741
742        if num_iterations.is_one() {
743            return true;
744        }
745
746        let elements_per_iteration: TDim = a_shape
747            .iter()
748            .zip(b_shape.iter())
749            .skip_while(|(_, b_dim)| **b_dim == 1.to_dim())
750            .map(|(_, b_dim)| b_dim)
751            .product();
752
753        if let Ok(num_element) = elements_per_iteration.to_i64() {
754            let required_alignment = vector_size();
755            (num_element as usize).is_multiple_of(required_alignment)
756        } else {
757            false
758        }
759    }
760    fn check_input_shapes(a_shape: &[TDim], b_shape: &[TDim]) -> bool {
761        if a_shape.len() != b_shape.len() {
762            return false;
763        };
764
765        let unicast_possible = a_shape
766            .iter()
767            .zip(b_shape.iter())
768            .skip_while(|(_, b_dim)| **b_dim == 1.to_dim())
769            .all(|(a_dim, b_dim)| a_dim == b_dim);
770        let unicast_is_aligned = Self::check_b_alignement(a_shape, b_shape);
771
772        unicast_possible && unicast_is_aligned
773    }
774}
775
776impl PartialEq for OptBinUnicast {
777    fn eq(&self, other: &Self) -> bool {
778        *self.binop == *other.binop
779    }
780}
781impl Eq for OptBinUnicast {}
782
783impl Op for OptBinUnicast {
784    fn name(&self) -> StaticName {
785        format!("Opt{}Unicast", self.binop.name()).into()
786    }
787
788    op_as_typed_op!();
789}
790
791impl EvalOp for OptBinUnicast {
792    fn is_stateless(&self) -> bool {
793        true
794    }
795
796    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
797        let (a, b) = args_2!(inputs);
798        // The unicast fast path indexes each input's storage via at_prefix +
799        // as_slice_mut, which uses `strides[i-1]` to size the resulting slice
800        // (data/src/tensor/view.rs:99). That formula only matches ∏(shape[i..])
801        // when the tensor has natural C-order strides. Producers like
802        // Tensor::insert_axis leave non-natural strides on a tensor (e.g.
803        // shape `[1, 1, 640]` with strides `[1, 1, 1]` after two insert_axis
804        // on a `[640]` tensor), which silently breaks the slice math. Fall
805        // back to the generic broadcasting eval when either operand is not in
806        // natural strides (or has a storage size that doesn't match the
807        // declared shape).
808        let a_natural = a.len() == a.shape().iter().product::<usize>()
809            && a.strides() == &*Tensor::natural_strides(a.shape());
810        let b_natural = b.len() == b.shape().iter().product::<usize>()
811            && b.strides() == &*Tensor::natural_strides(b.shape());
812        if !a_natural || !b_natural {
813            let c_dt = self.binop.result_datum_type(a.datum_type(), b.datum_type())?;
814            return Ok(tvec!(self.binop.eval(a, b, c_dt)?.into_tvalue()));
815        }
816
817        let mut a = a.into_tensor();
818        // b's leading unary axes are the ones a repeats over; past them b matches
819        // a exactly (check_input_shapes), so one kernel call covers b.len()
820        // elements of a and b is consumed in lockstep within each.
821        debug_assert!(
822            b.shape().iter().zip(a.shape()).skip_while(|(b, _)| **b == 1).all(|(b, a)| b == a),
823            "unicast b {:?} does not line up with a {:?}",
824            b.shape(),
825            a.shape()
826        );
827        let period = b.len();
828        tract_linalg::multithread::par_bin(&*self.eval_fn, &mut a, &b, period, BShare::Lockstep)?;
829
830        Ok(tvec!(a.into_tvalue()))
831    }
832}
833
834impl TypedOp for OptBinUnicast {
835    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
836        ensure!(Self::check_input_shapes(&inputs[0].shape, &inputs[1].shape));
837        let out_dt = self.binop.result_datum_type(inputs[0].datum_type, inputs[1].datum_type)?;
838        let out_shape = inputs[0].shape.clone();
839        Ok(tvec!(out_dt.fact(out_shape)))
840    }
841
842    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
843        let count: TDim = self.output_facts(inputs)?[0].shape.iter().product();
844        Ok(self
845            .binop
846            .cost_per_element(inputs[0].datum_type)
847            .into_iter()
848            .map(|(c, n)| (c, count.clone() * n))
849            .collect())
850    }
851
852    as_op!();
853}
854
855#[macro_export]
856macro_rules! bin_to_super_type {
857    ($func:ident, $Op:ident,
858     $(codegen: $codegen:expr,)?
859     $(cost: $cost:expr,)?
860     $(declutter: $declutter:expr,)?
861     $(eval_in_a: $eval_in_a:expr,)?
862     $(eval_override: $eval_override: expr,)?
863     $(linalg: $linalg:ident,)?
864     $(operating_datum_type: $operating_datum_type:expr,)?
865     $(is_commutative: $is_commutative:expr,)?
866     $(neutral_element: $neutral_element:expr,)?
867     $(absorbing_element: $absorbing_element:expr,)?
868     $(out_of_place: $out_of_place:expr,)?
869     $(validation: $validation:expr,)?
870     $(q: $([$($typ_dt:ident),*] => $cab_dt:expr),* ;)?
871     $(q_op_on_f32: $q_op_on_f32:expr,)?
872     $( [$($typ:ident),*] => $cab:expr),*) => {
873        #[derive(Debug, Clone, Hash, PartialEq, Eq)]
874        pub struct $Op;
875        #[allow(clippy::redundant_closure_call)]
876        impl $crate::ops::binary::BinMiniOp for $Op {
877            fn name(&self) -> &'static str {
878                stringify!($Op)
879            }
880
881            fn eval_out_of_place(&self, c: &mut Tensor, a: &Tensor, b: &Tensor) -> TractResult<()> {
882                $(if $out_of_place(c, a, b)? { return Ok(()) } )?
883                    // Same-shape fast path: bypass ndarray Zip when c, a, b
884                    // share the same shape (and hence same len for plain
885                    // storage). Iterate over slices directly.
886                    if c.shape() == a.shape() && a.shape() == b.shape() {
887                        $(
888                            $(if c.datum_type() == $typ::datum_type() {
889                                let cab: fn(&mut $typ, &$typ, &$typ) -> () = $cab;
890                                let a_plain = a.try_as_plain()?;
891                                let a_slice = a_plain.as_slice::<$typ>()?;
892                                let b_plain = b.try_as_plain()?;
893                                let b_slice = b_plain.as_slice::<$typ>()?;
894                                let mut c_plain = c.try_as_plain_mut()?;
895                                let c_slice = c_plain.as_slice_mut::<$typ>()?;
896                                debug_assert_eq!(c_slice.len(), a_slice.len());
897                                debug_assert_eq!(c_slice.len(), b_slice.len());
898                                let len = c_slice.len();
899                                tract_linalg::multithread::par_chunks_mut(c_slice, 1, len, |first_row, c_chunk| {
900                                    let n = c_chunk.len();
901                                    let a_chunk = &a_slice[first_row..first_row + n];
902                                    let b_chunk = &b_slice[first_row..first_row + n];
903                                    for ((cv, av), bv) in c_chunk.iter_mut().zip(a_chunk.iter()).zip(b_chunk.iter()) {
904                                        cab(cv, av, bv);
905                                    }
906                                    Ok(())
907                                })?;
908                                return Ok(())
909                            })*
910                        )*
911                        $(
912                            $(
913                                $(if a.datum_type().unquantized() == <$typ_dt>::datum_type().unquantized() {
914                                    let cab: fn(&mut $typ_dt, &$typ_dt, &$typ_dt, i32, f32) -> () = $cab_dt;
915                                    let (zp, scale) = a.datum_type().qparams().map(|q| q.zp_scale()).unwrap_or((0, 1.));
916                                    let a_plain = a.try_as_plain()?;
917                                    let a_slice = a_plain.as_slice::<$typ_dt>()?;
918                                    let b_plain = b.try_as_plain()?;
919                                    let b_slice = b_plain.as_slice::<$typ_dt>()?;
920                                    let mut c_plain = c.try_as_plain_mut()?;
921                                    let c_slice = c_plain.as_slice_mut::<$typ_dt>()?;
922                                    for ((cv, av), bv) in c_slice.iter_mut().zip(a_slice.iter()).zip(b_slice.iter()) {
923                                        cab(cv, av, bv, zp, scale);
924                                    }
925                                    return Ok(())
926                                })*
927                            )*
928                        )?
929                    }
930                    $(
931                        $(if c.datum_type() == $typ::datum_type() {
932                            let a = a.to_plain_array_view::<$typ>()?;
933                            let b = b.to_plain_array_view::<$typ>()?;
934                            let mut c_plain = c.try_as_plain_mut()?;
935                            let mut c = c_plain.to_array_view_mut::<$typ>()?;
936                            $crate::ndarray::Zip::from(&mut c).and_broadcast(a).and_broadcast(b).for_each($cab);
937                            return Ok(())
938                        })*
939                     )*
940                    $(
941                        $(
942                            $(if a.datum_type().unquantized() == <$typ_dt>::datum_type().unquantized() {
943                                let cab: fn(&mut $typ_dt, &$typ_dt, &$typ_dt, i32, f32) -> () = $cab_dt;
944                                let (zp, scale) = a.datum_type().qparams().map(|q| q.zp_scale()).unwrap_or((0, 1.));
945                                let a = a.to_plain_array_view::<$typ_dt>()?;
946                                let b = b.to_plain_array_view::<$typ_dt>()?;
947                                let mut c_plain = c.try_as_plain_mut()?;
948                                let mut c = c_plain.to_array_view_mut::<$typ_dt>()?;
949                                $crate::ndarray::Zip::from(&mut c).and_broadcast(a).and_broadcast(b).for_each(|c, a, b| cab(c, a, b, zp, scale));
950                                return Ok(())
951                            }
952                            )*
953                         )*
954                     )?
955                    bail!("{} does not support {:?} (out of place)", self.name(), c.datum_type());
956            }
957
958            $(fn is_commutative(&self) -> bool {
959                $is_commutative
960            })?
961            $(fn neutral_element(&self) -> Option<i64> {
962                Some($neutral_element)
963            })?
964            $(fn absorbing_element(&self) -> Option<i64> {
965                Some($absorbing_element)
966            })?
967            fn eval_in_a(&self, a: &mut Tensor, b: &Tensor) -> TractResult<()> {
968                // c and a are same type
969                $(if $eval_in_a(a, b)? { return Ok(()) } )?
970                // Same-shape fast path: bypass ndarray Zip when a and b share
971                // the same shape (and hence same len for plain storage).
972                if a.shape() == b.shape() {
973                    $(
974                        $(if b.datum_type() == $typ::datum_type() {
975                            let cab: fn(&mut $typ, &$typ, &$typ) -> () = $cab;
976                            let b_plain = b.try_as_plain()?;
977                            let b_slice = b_plain.as_slice::<$typ>()?;
978                            let mut a_plain = a.try_as_plain_mut()?;
979                            let a_slice = a_plain.as_slice_mut::<$typ>()?;
980                            debug_assert_eq!(a_slice.len(), b_slice.len());
981                            let len = a_slice.len();
982                            tract_linalg::multithread::par_chunks_mut(a_slice, 1, len, |first_row, a_chunk| {
983                                let n = a_chunk.len();
984                                let b_chunk = &b_slice[first_row..first_row + n];
985                                for (av, bv) in a_chunk.iter_mut().zip(b_chunk.iter()) {
986                                    cab(av, &av.clone(), bv);
987                                }
988                                Ok(())
989                            })?;
990                            return Ok(())
991                        })*
992                    )*
993                    $(
994                        $(
995                            $(if a.datum_type().unquantized() == <$typ_dt>::datum_type().unquantized() {
996                                let cab: fn(&mut $typ_dt, &$typ_dt, &$typ_dt, i32, f32) -> () = $cab_dt;
997                                let (zp, scale) = a.datum_type().qparams().map(|q| q.zp_scale()).unwrap_or((0, 1.));
998                                let b_plain = b.try_as_plain()?;
999                                let b_slice = b_plain.as_slice::<$typ_dt>()?;
1000                                let mut a_plain = a.try_as_plain_mut()?;
1001                                let a_slice = a_plain.as_slice_mut::<$typ_dt>()?;
1002                                for (av, bv) in a_slice.iter_mut().zip(b_slice.iter()) {
1003                                    cab(av, &(av.clone()), bv, zp, scale);
1004                                }
1005                                return Ok(())
1006                            })*
1007                        )*
1008                    )?
1009                }
1010                $(
1011                    $(if b.datum_type() == $typ::datum_type() {
1012                        let cab: fn(&mut $typ, &$typ, &$typ) -> () = $cab;
1013                        let b = b.to_plain_array_view::<$typ>()?;
1014                        let mut a_plain = a.try_as_plain_mut()?;
1015                        let mut a = a_plain.to_array_view_mut::<$typ>()?;
1016                        $crate::ndarray::Zip::from(&mut a).and_broadcast(b).for_each(|a, b| cab(a, &a.clone(), b));
1017                        return Ok(())
1018                    })*
1019                )*
1020                $(
1021                    $(
1022                        $(if a.datum_type().unquantized() == <$typ_dt>::datum_type().unquantized() {
1023                            let cab: fn(&mut $typ_dt, &$typ_dt, &$typ_dt, i32, f32) -> () = $cab_dt;
1024                            let (zp, scale) = a.datum_type().qparams().map(|q| q.zp_scale()).unwrap_or((0, 1.));
1025                            let mut a_plain = a.try_as_plain_mut()?;
1026                            let mut a = a_plain.to_array_view_mut::<$typ_dt>()?;
1027                            let b = b.to_plain_array_view::<$typ_dt>()?;
1028                            $crate::ndarray::Zip::from(&mut a).and_broadcast(b).for_each(|a, b| {
1029                                cab(a, &(a.clone()), b, zp, scale)
1030                            });
1031                            return Ok(())
1032                        })*
1033                    )*
1034                )?
1035                bail!("{} does not support {:?} (eval in a)", self.name(), a.datum_type());
1036            }
1037
1038            $(fn eval(&self, a: TValue, b: TValue, c_dt: DatumType) -> TractResult<Tensor> {
1039                $eval_override(a, b, c_dt)
1040            })?
1041
1042            fn result_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType> {
1043                if a.unquantized() == b.unquantized() {
1044                    if a.is_quantized() || !b.is_quantized() {
1045                        return Ok(a)
1046                    }
1047                    else {
1048                        return Ok(b)
1049                    }
1050                }
1051                self.operating_datum_type(a, b)
1052            }
1053
1054                $(
1055                    fn declutter(
1056                        &self,
1057                        model: &TypedModel,
1058                        node: &TypedNode,
1059                        ) -> TractResult<Option<TypedModelPatch>> {
1060                        ($declutter)(self, model, node)
1061                    }
1062                 )?
1063                $(
1064                    fn codegen(
1065                        &self,
1066                        model: &TypedModel,
1067                        node: &TypedNode,
1068                        a: &Arc<Tensor>,
1069                        ) -> TractResult<Option<TypedModelPatch>> {
1070                        ($codegen)(self, model, node, a)
1071                    }
1072                 )?
1073                $(
1074                    fn cost_per_element(&self, dt: DatumType) -> TVec<(Cost, usize)> {
1075                        ($cost)(dt)
1076                    }
1077                 )?
1078                $(
1079                    fn validation(&self) -> Validation {
1080                        $validation
1081                    }
1082                 )?
1083                $(
1084                    fn as_linalg_binop(&self) -> Option<tract_linalg::BinOp> {
1085                        Some(tract_linalg::BinOp::$linalg)
1086                    }
1087                 )?
1088                $(
1089                    fn operating_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType> {
1090                        ($operating_datum_type)(a, b)
1091                    })?
1092
1093
1094            /// Default simple binary operation for QFormat where
1095            /// we dequantise & apply requested operation in float & requantize it
1096            /// several implementation are provided with pro & con
1097            #[allow(unused_variables)]
1098            fn maybe_eval_qbinary_as_float_op(
1099                &self,
1100                a: &TValue,
1101                b: &TValue,
1102                c_dt: &DatumType,
1103            ) -> TractResult<Option<Tensor>> {
1104                $(
1105                    /// Implementation strive to minimise memory allocation and access
1106                    /// we apply only if type is QU8 zp_scale datum type
1107                    /// maybe more suited for large models tensors
1108                    fn memory_optimised_q_binary_as_float_op(
1109                        a: &TValue,
1110                        b: &TValue,
1111                        c_dt: &DatumType,
1112                    ) -> TractResult<Option<Tensor>> {
1113                        if let (DatumType::QU8(QParams::ZpScale {zero_point: a_zp, scale: a_scale}),
1114                                DatumType::QU8(QParams::ZpScale {zero_point: b_zp, scale: b_scale}),
1115                                DatumType::QU8(QParams::ZpScale {zero_point: c_zp, scale: c_scale})) =
1116                            (a.datum_type(), b.datum_type(), c_dt)
1117                        {
1118                            let c_inv_scale = 1.0 / c_scale;
1119                            let a = a.to_plain_array_view::<u8>()?;
1120                            let b = b.to_plain_array_view::<u8>()?;
1121                            let c_shape = $crate::broadcast::multi_broadcast(&[a.shape(), b.shape()])?;
1122                            let mut c = Tensor::zero_dt(*c_dt, &c_shape)?;
1123                            let mut c_plain = c.try_as_plain_mut()?;
1124                            let view = c_plain.to_array_view_mut::<u8>()?;
1125                            $crate::ndarray::Zip::from(view).and_broadcast(a).and_broadcast(b).for_each(|c, a, b| {
1126                                *c = (scale_by($q_op_on_f32(
1127                                            ((*a as i32 - a_zp as i32) as f32 * a_scale),
1128                                            ((*b as i32 - b_zp as i32) as f32 * b_scale),
1129                                ), c_inv_scale) as i32
1130                                    + *c_zp as i32)
1131                                    .clamp_cast()
1132                            });
1133                            return Ok(Some(c));
1134                        }
1135                        Ok(None)
1136                    }
1137
1138                    /// Apply to all Q types
1139                    /// Take more memory but hopefully faster than memory_optimised_q_binary_as_float_op
1140                    /// especially once cast_to_dt will have will have vectorized implementations
1141                    fn generic_q_binary_as_float_op(
1142                        a: &TValue,
1143                        b: &TValue,
1144                        c_dt: &DatumType,
1145                        accumulator_dt: DatumType
1146                    ) -> TractResult<Option<Tensor>> {
1147                        if a.datum_type().is_quantized() && b.datum_type().is_quantized() && c_dt.is_quantized() {
1148                            let a = a.cast_to_dt(accumulator_dt)?.into_owned();
1149                            let b = b.cast_to_dt(accumulator_dt)?.into_owned();
1150                            let c_shape = $crate::broadcast::multi_broadcast(&[a.shape(), b.shape()])?;
1151                            let mut c = Tensor::zero_dt(accumulator_dt, &c_shape)?;
1152                            match accumulator_dt {
1153                                DatumType::F32 => {
1154                                    let mut c_plain = c.try_as_plain_mut()?;
1155                                    let view = c_plain.to_array_view_mut::<f32>()?;
1156                                    $crate::ndarray::Zip::from(view).and_broadcast(a.try_as_plain()?.to_array_view()?).and_broadcast(b.try_as_plain()?.to_array_view()?).for_each(|c, a, b| {
1157                                        *c = $q_op_on_f32(*a,*b);
1158                                    })
1159                                },
1160                                other => bail!("unexpected accumulator data type as {:?}", other)
1161                            };
1162
1163                            return Ok(Some(c.cast_to_dt(*c_dt)?.into_owned()));
1164                        }
1165                        Ok(None)
1166                    }
1167
1168                    if let Some(c) = memory_optimised_q_binary_as_float_op(a, b, c_dt)? {
1169                        return Ok(Some(c));
1170                    }
1171                    if let Some(d) = generic_q_binary_as_float_op(a, b, c_dt, DatumType::F32)? {
1172                        return Ok(Some(d));
1173                    }
1174                )?
1175                Ok(None)
1176            }
1177        }
1178
1179        pub fn $func() -> $crate::ops::binary::TypedBinOp {
1180            $crate::ops::binary::TypedBinOp(Box::new($Op), None)
1181        }
1182    };
1183}
1184
1185#[derive(Debug)]
1186pub(crate) struct OneUniformInput {
1187    pub uni: Arc<Tensor>,
1188    pub var: OutletId,
1189    pub left_is_uniform: bool,
1190}
1191
1192pub(crate) fn one_input_is_uniform(
1193    model: &TypedModel,
1194    node: &TypedNode,
1195) -> TractResult<Option<OneUniformInput>> {
1196    if let &[a, b] = &*model.node_input_facts(node.id)? {
1197        let uni = if let Some(a) = &a.uniform {
1198            OneUniformInput { uni: a.clone(), var: node.inputs[1], left_is_uniform: true }
1199        } else if let Some(b) = &b.uniform {
1200            OneUniformInput { uni: b.clone(), var: node.inputs[0], left_is_uniform: false }
1201        } else {
1202            return Ok(None);
1203        };
1204        let var_fact = [a, b][uni.left_is_uniform as usize];
1205        let uni_fact = [a, b][!uni.left_is_uniform as usize];
1206        if izip!(var_fact.shape.iter(), uni_fact.shape.iter()).all(|(v, u)| u.is_one() || u == v) {
1207            return Ok(Some(uni));
1208        }
1209    }
1210    Ok(None)
1211}
1212
1213#[cfg(test)]
1214mod tests {
1215    use super::*;
1216
1217    /// Reproducer for the OptBinUnicast panic seen on Nemotron decoder CI
1218    /// (cuda-lovelace + Darwin). A 1-D tensor that goes through `insert_axis`
1219    /// twice ends up with declared shape `[1, 1, 640]` but strides `[1, 1, 1]`
1220    /// instead of the natural `[640, 640, 1]`. TensorView::at_prefix then
1221    /// returns a view whose `len()` reads `strides[1] = 1`, so the unicast
1222    /// kernel sees `a.len = 1, b.len = 640` and OOBs into the tile buffer.
1223    ///
1224    /// Pre-fix this test panics inside `linalg/src/frame/unicast.rs` with
1225    /// "range end index 640 out of range for slice of length …". With the
1226    /// natural-strides guard in `OptBinUnicast::eval`, the call falls back to
1227    /// `BinMiniOp::eval` and produces correct output.
1228    #[test]
1229    fn opt_bin_unicast_falls_back_on_non_natural_strides() {
1230        // Construct `a` the way the LSTM bias path does: build a 640-element
1231        // 1-D tensor, then insert two leading unit dims.
1232        let a_data: Vec<f32> = (0..640).map(|i| i as f32).collect();
1233        let mut a = tensor1(&a_data);
1234        a.insert_axis(0).unwrap();
1235        a.insert_axis(0).unwrap();
1236        assert_eq!(a.shape(), &[1, 1, 640]);
1237        assert_eq!(a.strides(), &[1, 1, 1]);
1238        assert_ne!(a.strides(), &*Tensor::natural_strides(a.shape()));
1239
1240        // `b` is a normal contiguous tensor of the same declared shape.
1241        let b_data: Vec<f32> = vec![1.0; 640];
1242        let mut b = tensor1(&b_data);
1243        b.insert_axis(0).unwrap();
1244        b.insert_axis(0).unwrap();
1245        // Reset b to natural strides so we exercise only the a-broken path
1246        // and let the b-side go through cleanly.
1247        b = b.into_shape(&[1, 1, 640]).unwrap();
1248
1249        let linalg_fn = tract_linalg::bin_unicast(f32::datum_type(), BinOp::Add)
1250            .expect("f32 unicast Add kernel available");
1251        let op = OptBinUnicast { binop: Box::new(Add), eval_fn: Arc::from(linalg_fn) };
1252
1253        let out = op.eval(tvec!(a.into_tvalue(), b.into_tvalue())).unwrap();
1254        let out = &out[0];
1255        assert_eq!(out.shape(), &[1, 1, 640]);
1256        let plain = out.try_as_plain().unwrap();
1257        let out_slice = plain.as_slice::<f32>().unwrap();
1258        for (i, v) in out_slice.iter().enumerate() {
1259            assert_eq!(*v, i as f32 + 1.0, "mismatch at {i}");
1260        }
1261    }
1262
1263    /// A zero-sized outer dim (a symbolic batch or sequence resolving to 0) has
1264    /// no elements to write, so both fast paths must no-op. The kernels cannot
1265    /// be handed the empty tensor: `as_slice_mut` builds a slice from the null
1266    /// data pointer, and the by-scalar kernel reads `b[0]` off a zero-length
1267    /// slice.
1268    #[test]
1269    fn zero_sized_outer_dim_is_a_noop() {
1270        let a = Tensor::zero::<f32>(&[0, 4, 8]).unwrap();
1271        let b = Tensor::zero::<f32>(&[0, 4, 1]).unwrap();
1272        let linalg_fn = tract_linalg::bin_by_scalar(f32::datum_type(), BinOp::Add)
1273            .expect("f32 by_scalar Add kernel available");
1274        let op = OptBinByScalar { binop: Box::new(Add), eval_fn: Arc::from(linalg_fn) };
1275        let out = op.eval(tvec!(a.into_tvalue(), b.into_tvalue())).unwrap();
1276        assert_eq!(out[0].shape(), &[0, 4, 8]);
1277
1278        let a = Tensor::zero::<f32>(&[0, 4, 16]).unwrap();
1279        let b = Tensor::zero::<f32>(&[1, 4, 16]).unwrap();
1280        let linalg_fn = tract_linalg::bin_unicast(f32::datum_type(), BinOp::Add)
1281            .expect("f32 unicast Add kernel available");
1282        let op = OptBinUnicast { binop: Box::new(Add), eval_fn: Arc::from(linalg_fn) };
1283        let out = op.eval(tvec!(a.into_tvalue(), b.into_tvalue())).unwrap();
1284        assert_eq!(out[0].shape(), &[0, 4, 16]);
1285    }
1286
1287    /// A quantized uniform input that dequantizes to 0 is neutral for Add on
1288    /// either side, but the node still re-encodes its variable input into the
1289    /// output quantization parameters. `declutter_neutral` must therefore cast
1290    /// the variable input, which is inlet 1 here, not inlet 0.
1291    #[test]
1292    fn q_neutral_add_requantizes_the_variable_input() -> TractResult<()> {
1293        let x_dt = DatumType::QU8(QParams::ZpScale { zero_point: 10, scale: 0.02 });
1294        let zero_dt = DatumType::QU8(QParams::ZpScale { zero_point: 61, scale: 1. });
1295        let out_dt = DatumType::QU8(QParams::ZpScale { zero_point: 0, scale: 0.5 });
1296
1297        let mut model = TypedModel::default();
1298        let x = model.add_source("x", x_dt.fact([4]))?;
1299        // A quantized zero tensor is filled with its zero point, so this
1300        // constant dequantizes to 0.0 everywhere.
1301        let zero = model.add_const("zero", Tensor::zero_dt(zero_dt, &[4])?)?;
1302        let add = model.wire_node("add", TypedBinOp(Box::new(Add), Some(out_dt)), &[zero, x])?[0];
1303        model.select_output_outlets(&[add])?;
1304
1305        let mut input = Tensor::zero_dt(x_dt, &[4])?;
1306        input.try_as_plain_mut()?.as_slice_mut::<u8>()?.copy_from_slice(&[10, 35, 60, 200]);
1307        let input = tvec!(input.into_tvalue());
1308
1309        let before = model.clone().into_runnable()?.run(input.clone())?;
1310        let decluttered = model.into_decluttered()?;
1311        assert!(decluttered.nodes().iter().all(|n| n.op_as::<TypedBinOp>().is_none()));
1312        let after = decluttered.into_runnable()?.run(input)?;
1313        assert_eq!(&*before[0], &*after[0]);
1314        Ok(())
1315    }
1316}