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 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 #[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 #[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 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 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 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 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 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 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 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 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 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 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
476fn 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 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
514fn 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 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 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 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 let n_blocks: usize = a.shape()[..first_unary_axis].iter().product();
693 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 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 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 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 $(if $eval_in_a(a, b)? { return Ok(()) } )?
970 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 #[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 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 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 #[test]
1229 fn opt_bin_unicast_falls_back_on_non_natural_strides() {
1230 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 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 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 #[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 #[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 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}