tract_core/ops/
quant.rs

1#![allow(clippy::unnecessary_cast)]
2
3use crate::internal::*;
4use crate::ops::element_wise::ElementWiseOp;
5use crate::ops::math::QScale;
6use num_traits::AsPrimitive;
7use tract_linalg::Scaler;
8use tract_linalg::lut::Lut;
9use tract_linalg::mmm::RoundingPolicy;
10
11use super::binary::TypedBinOp;
12use super::math::round_ties_to_even;
13
14pub fn quantize_linear_f32_u8(x: f32, scale: f32, zero_point: i32) -> u8 {
15    (((x * scale).round() as i32) + zero_point).clamp(u8::MIN as i32, u8::MAX as i32) as u8
16}
17
18pub fn quantize_linear_f32_i8(x: f32, scale: f32, zero_point: i32) -> i8 {
19    (((x * scale).round() as i32) + zero_point).clamp(i8::MIN as i32, i8::MAX as i32) as i8
20}
21
22element_wise_oop!(quantize_linear_u8,
23 QuantizeLinearU8 {
24     scale: f32,
25     zero_point: u8
26 },
27 [f16] => u8 |op, xs, ys| {
28     xs.iter().zip(ys.iter_mut()).for_each(|(x,y)|
29                                           *y = quantize_linear_f32_u8(x.to_f32(), op.scale, op.zero_point as i32)
30                                          );
31     Ok(())
32 },
33 [f32,i32] => u8 |op, xs, ys| {
34     xs.iter().zip(ys.iter_mut()).for_each(|(x,y)|
35                                           *y = quantize_linear_f32_u8(*x as f32, op.scale, op.zero_point as i32)
36                                          );
37     Ok(())
38 };
39 info: info_quantize_linear_u8
40);
41
42fn info_quantize_linear_u8(q: &QuantizeLinearU8) -> TractResult<Vec<String>> {
43    Ok(vec![format!(
44        "scale: {} zero_point: {} 1/scale: {}",
45        q.scale,
46        q.zero_point,
47        q.scale.recip()
48    )])
49}
50
51element_wise_oop!(quantize_linear_i8,
52 QuantizeLinearI8 {
53     scale: f32,
54     zero_point: i8
55 },
56 [f32,i32] => i8 |op, xs, ys| {
57     xs.iter().zip(ys.iter_mut()).for_each(|(x,y)|
58                                           *y = quantize_linear_f32_i8(*x as f32, op.scale, op.zero_point as i32)
59                                          );
60     Ok(())
61 };
62 info: info_quantize_linear_i8
63);
64
65fn info_quantize_linear_i8(q: &QuantizeLinearI8) -> TractResult<Vec<String>> {
66    Ok(vec![format!(
67        "scale: {} zero_point: {} 1/scale: {}",
68        q.scale,
69        q.zero_point,
70        q.scale.recip()
71    )])
72}
73
74#[derive(Clone, Debug, new, PartialEq)]
75pub struct DequantizeLinearF32 {
76    pub scale: f32,
77    pub zero_point: i32,
78}
79
80impl Eq for DequantizeLinearF32 {}
81
82impl DequantizeLinearF32 {
83    fn eval_t<T: Datum + AsPrimitive<i32>>(&self, input: &Tensor) -> TractResult<Tensor> {
84        let mut output = unsafe { Tensor::uninitialized::<f32>(input.shape())? };
85        input
86            .try_as_plain()?
87            .as_slice::<T>()?
88            .iter()
89            .zip(output.try_as_plain_mut()?.as_slice_mut::<f32>()?.iter_mut())
90            .for_each(|(x, y)| *y = (x.as_() - self.zero_point) as f32 * self.scale);
91        Ok(output)
92    }
93}
94
95impl Op for DequantizeLinearF32 {
96    fn name(&self) -> StaticName {
97        "DequantizeLinearF32".into()
98    }
99
100    fn info(&self) -> TractResult<Vec<String>> {
101        Ok(vec![format!("scale: {} zero_point: {}", self.scale, self.zero_point)])
102    }
103
104    fn validation(&self) -> Validation {
105        Validation::Accurate
106    }
107
108    op_as_typed_op!();
109}
110
111impl EvalOp for DequantizeLinearF32 {
112    op_out_of_plan!();
113    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
114        let output = match inputs[0].datum_type() {
115            DatumType::I8 => self.eval_t::<i8>(&inputs[0])?,
116            DatumType::I32 => self.eval_t::<i32>(&inputs[0])?,
117            DatumType::U8 => self.eval_t::<u8>(&inputs[0])?,
118            dt => bail!("Unsupported type {:?}", dt),
119        };
120        Ok(tvec!(output.into_tvalue()))
121    }
122}
123
124impl TypedOp for DequantizeLinearF32 {
125    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
126        Ok(tvec!(f32::datum_type().fact(inputs[0].shape.clone())))
127    }
128
129    fn axes_mapping(
130        &self,
131        inputs: &[&TypedFact],
132        outputs: &[&TypedFact],
133    ) -> TractResult<AxesMapping> {
134        AxesMapping::natural(inputs, outputs)
135    }
136
137    fn change_axes(
138        &self,
139        model: &TypedModel,
140        node: &TypedNode,
141        _io: InOut,
142        change: &AxisOp,
143    ) -> TractResult<Option<AxisChangeConsequence>> {
144        Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
145    }
146
147    fn declutter(
148        &self,
149        model: &TypedModel,
150        dequant: &TypedNode,
151    ) -> TractResult<Option<TypedModelPatch>> {
152        let mut current = dequant;
153        let incoming_dt = model.node_input_facts(dequant.id)?[0].datum_type;
154        while let Some(quant) = model.single_succ(current.id)? {
155            let q_params = if let Some(op) = quant.op_as::<ElementWiseOp>() {
156                if let Some(mop) = op.0.downcast_ref::<QuantizeLinearU8>() {
157                    Some((mop.scale, mop.zero_point as i32, u8::datum_type()))
158                } else {
159                    op.0.downcast_ref::<QuantizeLinearI8>()
160                        .map(|mop| (mop.scale, mop.zero_point as i32, i8::datum_type()))
161                }
162            } else {
163                None
164            };
165            if let Some((scale, zero_point, dt)) = q_params {
166                // first, try Op::quantize() on all ops in the chain
167                let mut patch = TypedModelPatch::default();
168                let mut wire: OutletId = patch.tap_model(model, dequant.inputs[0])?;
169                let mut next = model.single_succ(dequant.id)?.unwrap();
170                loop {
171                    if let Some(op) = next
172                        .op
173                        .quantize(model, dequant, dt, scale, zero_point)
174                        .with_context(|| format!("Quantizing {next}"))?
175                    {
176                        wire = patch.wire_node(&*next.name, op, [wire].as_ref())?[0];
177                    } else {
178                        break;
179                    }
180                    if next.id == current.id {
181                        patch.shunt_outside(model, OutletId::new(quant.id, 0), wire)?;
182                        return Ok(Some(patch));
183                    } else {
184                        next = model.single_succ(next.id)?.unwrap();
185                    }
186                }
187                // or else make a lookup table
188                if incoming_dt == DatumType::I8 || incoming_dt == DatumType::U8 {
189                    let mut adhoc_model = TypedModel::default();
190                    let mut wire = adhoc_model.add_source("ad-hoc", dt.fact([256]))?;
191                    let mut next = model.single_succ(dequant.id)?.unwrap();
192                    let mut name = None;
193                    // plug in dequant
194                    wire = adhoc_model.wire_node(
195                        &*dequant.name,
196                        dequant.op.clone(),
197                        [wire].as_ref(),
198                    )?[0];
199                    while next.id != quant.id {
200                        name.get_or_insert(&*next.name);
201                        wire =
202                            adhoc_model.wire_node(&*next.name, next.op.clone(), [wire].as_ref())?
203                                [0];
204                        next = model.single_succ(next.id)?.unwrap();
205                    }
206                    // plug in quant
207                    wire =
208                        adhoc_model.wire_node(&*quant.name, quant.op.clone(), [wire].as_ref())?[0];
209                    adhoc_model.select_output_outlets(&[wire])?;
210                    let input = (0u8..=255).collect::<Vec<u8>>();
211                    let input = match dt {
212                        DatumType::I8 => unsafe {
213                            tensor1(std::mem::transmute::<&[u8], &[i8]>(&*input))
214                        },
215                        DatumType::U8 => tensor1(&input),
216                        _ => unreachable!(),
217                    };
218                    let output =
219                        SimplePlan::new(adhoc_model)?.run(tvec!(input.into_tvalue()))?.remove(0);
220                    let table: &[u8] = match dt {
221                        DatumType::I8 => unsafe {
222                            std::mem::transmute::<&[i8], &[u8]>(
223                                output.try_as_plain()?.as_slice::<i8>()?,
224                            )
225                        },
226                        DatumType::U8 => output.try_as_plain()?.as_slice::<u8>()?,
227                        _ => unreachable!(),
228                    };
229                    let op = lookup_table(tract_linalg::routines::lut_u8(table)?);
230                    let mut patch = TypedModelPatch::default();
231                    let mut wire: OutletId = patch.tap_model(model, dequant.inputs[0])?;
232
233                    wire = patch.wire_node(name.unwrap_or(&*dequant.name), op, [wire].as_ref())?[0];
234                    patch.shunt_outside(model, OutletId::new(quant.id, 0), wire)?;
235                    return Ok(Some(patch));
236                }
237            }
238            let (input_facts, output_facts) = model.node_facts(quant.id)?;
239            let invariants = quant
240                .op
241                .axes_mapping(&input_facts, &output_facts)
242                .with_context(|| format!("Querying invariants for {quant}"))?;
243            if invariants.is_element_wise_unary() {
244                current = quant;
245            } else {
246                break;
247            }
248        }
249        Ok(None)
250    }
251
252    as_op!();
253}
254
255element_wise_oop!(lookup_table,
256 LookupTable {
257     table: Box<dyn Lut>
258 },
259 [i8] => i8 |op, xs, ys| {
260     ys.copy_from_slice(xs);
261     unsafe {
262         let casted = std::slice::from_raw_parts_mut(ys.as_mut_ptr() as *mut u8, ys.len());
263         op.table.run(casted);
264     }
265     Ok(())
266 },
267 [u8] => u8 |op, xs, ys| {
268     ys.copy_from_slice(xs);
269     op.table.run(ys);
270     Ok(())
271 }
272);
273
274#[derive(Debug, Clone, Hash, PartialEq, Eq)]
275pub struct Scale;
276
277impl crate::ops::binary::BinMiniOp for Scale {
278    fn name(&self) -> &'static str {
279        "Scale"
280    }
281    fn result_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType> {
282        if !a.is_float() {
283            bail!("Scale left operand must be float, got {:?}", a);
284        }
285        Ok(b)
286    }
287
288    fn operating_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType> {
289        if !a.is_float() {
290            bail!("Scale left operand must be float, got {:?}", a);
291        }
292        Ok(b)
293    }
294
295    fn eval_out_of_place(&self, c: &mut Tensor, a: &Tensor, b: &Tensor) -> TractResult<()> {
296        let a = a.cast_to::<f32>()?;
297        let a = a.to_plain_array_view::<f32>()?;
298        unsafe fn eval_out_of_place_t<T: Datum + AsPrimitive<f32>>(
299            c: &mut Tensor,
300            a: &ndarray::ArrayViewD<f32>,
301            b: &Tensor,
302        ) where
303            f32: AsPrimitive<T>,
304        {
305            let b = unsafe { b.to_array_view_unchecked::<T>() };
306            let mut c = unsafe { c.to_array_view_mut_unchecked::<T>() };
307            ndarray::Zip::from(&mut c)
308                .and_broadcast(a)
309                .and_broadcast(b)
310                .for_each(|c, a, b| *c = scale_by(*b, *a))
311        }
312        unsafe { dispatch_numbers!(eval_out_of_place_t(b.datum_type())(c, &a, b)) }
313        Ok(())
314    }
315
316    fn eval_in_a(&self, a: &mut Tensor, b: &Tensor) -> TractResult<()> {
317        let mut a_plain = a.try_as_plain_mut()?;
318        let a = a_plain.to_array_view_mut::<f32>()?;
319        let b = b.to_plain_array_view::<f32>()?;
320        ndarray::Zip::from(a).and_broadcast(b).for_each(|a, b| *a = scale_by(*b, *a));
321        Ok(())
322    }
323
324    fn is_commutative(&self) -> bool {
325        false
326    }
327
328    fn declutter(
329        &self,
330        model: &TypedModel,
331        node: &TypedNode,
332    ) -> TractResult<Option<TypedModelPatch>> {
333        let a = model.outlet_fact(node.inputs[0])?;
334        if let Some(a) = &a.uniform {
335            if a.cast_to_scalar::<f32>()? == 1. {
336                return Ok(Some(TypedModelPatch::rewire(
337                    model,
338                    &node.inputs[1..2],
339                    &[node.id.into()],
340                    &|_p, x| Ok(x.into()),
341                )?));
342            } else if node.outputs[0].fact.datum_type == DatumType::I32 {
343                let factor = a.cast_to_scalar::<f32>()?;
344                let scaler = Scaler::new(factor, RoundingPolicy::Even);
345
346                let op = ElementWiseOp(Box::new(QScale { scaler }), None);
347                let patch =
348                    TypedModelPatch::replace_single_op(model, node, &node.inputs[1..2], op)?;
349
350                return Ok(Some(patch));
351            }
352        }
353        Ok(None)
354    }
355}
356
357#[inline]
358pub(crate) fn scale_by<T: Datum + AsPrimitive<f32>>(b: T, a: f32) -> T
359where
360    f32: AsPrimitive<T>,
361{
362    let b = b.as_();
363    (round_ties_to_even(b.abs() * a) * b.signum()).as_()
364}
365
366pub fn scale() -> TypedBinOp {
367    TypedBinOp(Box::new(Scale), None)
368}
369
370/// Offsets i8 integers as u8 integers.
371pub(crate) fn offset_i8_as_u8_elementwise(x: i8) -> u8 {
372    (x as u8).wrapping_add(128)
373}
374
375#[derive(Debug, Clone, PartialEq, Eq)]
376pub struct OffsetI8asU8;
377impl ElementWiseMiniOp for OffsetI8asU8 {
378    fn name(&self) -> String {
379        format!("{}{}", self.prefix(), stringify!(OffsetI8asU8))
380    }
381    fn output_type(&self, input_type: DatumType) -> Option<DatumType> {
382        Some(if let DatumType::QI8(qp) = input_type {
383            let (zp, scale) = qp.zp_scale();
384            DatumType::QU8(QParams::ZpScale { zero_point: zp + 128, scale })
385        } else if input_type == DatumType::I8 {
386            DatumType::U8
387        } else {
388            input_type
389        })
390    }
391    fn eval_out_of_place(&self, t: &Tensor, out_dt: Option<DatumType>) -> TractResult<Tensor> {
392        let output_type = out_dt.unwrap_or(self.output_type(t.datum_type()).unwrap());
393        let mut dst = unsafe { Tensor::uninitialized_dt(output_type, t.shape())? };
394        if t.datum_type().unquantized() == i8::datum_type() {
395            t.try_as_plain()?
396                .as_slice::<i8>()?
397                .iter()
398                .zip(dst.try_as_plain_mut()?.as_slice_mut::<u8>()?.iter_mut())
399                .for_each(|(x, y)| *y = offset_i8_as_u8_elementwise(*x));
400            return Ok(dst);
401        }
402
403        bail!("{} does not support {:?}", self.name(), t.datum_type());
404    }
405}
406
407pub fn offset_i8_as_u8() -> ElementWiseOp {
408    ElementWiseOp(Box::new(OffsetI8asU8 {}), None)
409}
410
411/// Offsets u8 integers as i8 integers.
412pub(crate) fn offset_u8_as_i8_elementwise(x: u8) -> i8 {
413    x.wrapping_sub(128) as i8
414}
415
416#[derive(Debug, Clone, PartialEq, Eq)]
417pub struct OffsetU8asI8;
418impl ElementWiseMiniOp for OffsetU8asI8 {
419    fn name(&self) -> String {
420        format!("{}{}", self.prefix(), stringify!(OffsetU8asI8))
421    }
422    fn output_type(&self, input_type: DatumType) -> Option<DatumType> {
423        Some(if let DatumType::QU8(qp) = input_type {
424            let (zp, scale) = qp.zp_scale();
425            DatumType::QI8(QParams::ZpScale { zero_point: zp - 128, scale })
426        } else if input_type == DatumType::U8 {
427            DatumType::I8
428        } else {
429            input_type
430        })
431    }
432    fn eval_out_of_place(&self, t: &Tensor, out_dt: Option<DatumType>) -> TractResult<Tensor> {
433        let output_type = out_dt.unwrap_or(self.output_type(t.datum_type()).unwrap());
434        let mut dst = unsafe { Tensor::uninitialized_dt(output_type, t.shape())? };
435        if t.datum_type().unquantized() == u8::datum_type() {
436            t.try_as_plain()?
437                .as_slice::<u8>()?
438                .iter()
439                .zip(dst.try_as_plain_mut()?.as_slice_mut::<i8>()?.iter_mut())
440                .for_each(|(x, y)| *y = offset_u8_as_i8_elementwise(*x));
441            return Ok(dst);
442        }
443
444        bail!("{} does not support {:?}", self.name(), t.datum_type());
445    }
446}
447pub fn offset_u8_as_i8() -> ElementWiseOp {
448    ElementWiseOp(Box::new(OffsetU8asI8 {}), None)
449}
450
451#[cfg(test)]
452pub mod scale {
453    use crate::internal::*;
454    use crate::ops::einsum::EinSum;
455    use crate::ops::math::round_ties_to_even;
456    use proptest::prelude::*;
457
458    fn test_scale(a: i8, b: i8, scale: f32) {
459        let expected = (((a as i32) * (b as i32)) as f32) / scale;
460        let expected = round_ties_to_even(expected.abs()) * expected.signum();
461        let expected = (expected as i32).clamp(-128, 127);
462        let expected = tensor2(&[[expected as i8]]);
463
464        let input = tvec!(tensor2(&[[b]]).into_tvalue());
465        let mut model = TypedModel::default();
466        let a = model.add_const("a", tensor2(&[[a]])).unwrap();
467        let b = model.add_source("b", i8::fact([1, 1])).unwrap();
468        let bias = model.add_const("bias", tensor0(0i32)).unwrap();
469        let a0 = model.add_const("a0", tensor0(0i8)).unwrap();
470        let a_scale = model.add_const("a_scale", tensor0(1f32)).unwrap();
471        let b0 = model.add_const("b0", tensor0(0i8)).unwrap();
472        let b_scale = model.add_const("b_scale", tensor0(1f32)).unwrap();
473        let c0 = model.add_const("c0", tensor0(0i8)).unwrap();
474        let c_scale = model.add_const("c_scale", tensor0(scale)).unwrap();
475        let op = EinSum {
476            axes: "mk,kn,,,,,,,->mn".parse().unwrap(),
477            operating_dt: i32::datum_type(),
478            q_params: Some(i8::datum_type()),
479        };
480        let output = model
481            .wire_node("mmm", op, &[a, b, bias, a0, a_scale, b0, b_scale, c0, c_scale])
482            .unwrap();
483        model.select_output_outlets(&output).unwrap();
484
485        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
486        assert_eq!(*plain[0], expected);
487
488        let optim = model.into_optimized().unwrap().into_runnable().unwrap().run(input).unwrap();
489        assert_eq!(*optim[0], expected);
490    }
491
492    proptest! {
493        #[test]
494        fn prop(a in any::<i8>(), b in any::<i8>(), scale in 0.00001f32..1000.) {
495            test_scale(a, b, scale);
496        }
497    }
498
499    #[test]
500    fn t1() {
501        test_scale(-117, 15, 37.753822);
502    }
503
504    #[test]
505    fn t2() {
506        test_scale(-4, -60, 475.21674);
507    }
508}
tract_core/ops/quant.rs

tract_core/ops/
quant.rs
