tract_core/ops/
cast.rs

1use crate::internal::*;
2use crate::ops::array::MultiBroadcastTo;
3
4pub fn cast(to: DatumType) -> Cast {
5    Cast { to }
6}
7
8pub fn wire_cast(
9    prefix: impl AsRef<str>,
10    target: &mut TypedModel,
11    inputs: &[OutletId],
12    operating_datum_type: DatumType,
13) -> TractResult<TVec<OutletId>> {
14    let prefix = prefix.as_ref();
15    let mut wires = tvec!();
16    for mut wire in inputs.iter().copied() {
17        if target.outlet_fact(wire)?.datum_type != operating_datum_type {
18            wire = target.wire_node(
19                target.unique_name(format!("{prefix}.cast")),
20                crate::ops::cast::cast(operating_datum_type),
21                &[wire],
22            )?[0];
23        }
24        wires.push(wire);
25    }
26    Ok(wires)
27}
28
29#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
30pub struct Cast {
31    pub to: DatumType,
32}
33
34impl Op for Cast {
35    fn name(&self) -> StaticName {
36        "Cast".into()
37    }
38
39    op_as_typed_op!();
40}
41
42impl EvalOp for Cast {
43    op_out_of_plan!();
44
45    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
46        let input = args_1!(inputs);
47        if input.datum_type() == self.to {
48            Ok(tvec!(input))
49        } else if input.datum_type() == TDim::datum_type() {
50            let mut tmp = Tensor::zero_dt(i64::datum_type(), input.shape())?;
51            let input_plain = input.try_as_plain()?;
52            let mut tmp_plain = tmp.try_as_plain_mut()?;
53            for (dim, i) in tract_itertools::izip!(
54                input_plain.as_slice::<TDim>()?,
55                tmp_plain.as_slice_mut::<i64>()?
56            ) {
57                *i = dim.eval(ctx.symbols).to_i64()?
58            }
59            Ok(tvec!(tmp.cast_to_dt(self.to)?.into_owned().into_tvalue()))
60        } else {
61            Ok(tvec!(input.cast_to_dt(self.to)?.into_owned().into_tvalue()))
62        }
63    }
64}
65
66impl TypedOp for Cast {
67    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
68        let mut fact = self.to.fact(inputs[0].shape.clone());
69        fact.uniform_tdim = inputs[0].uniform_tdim.clone();
70        if let Some(u) = &inputs[0].uniform
71            && let Ok(cast_u) = u.cast_to_dt(self.to)
72        {
73            fact.uniform = Some(std::sync::Arc::new(cast_u.into_owned()));
74        }
75        Ok(tvec!(fact))
76    }
77
78    fn input_roi(
79        &self,
80        model: &TypedModel,
81        node: &TypedNode,
82    ) -> TractResult<Option<TVec<Option<TDim>>>> {
83        crate::optim::propagate_roi::bubble_roi(model, node)
84    }
85
86    fn declutter(
87        &self,
88        model: &TypedModel,
89        node: &TypedNode,
90    ) -> TractResult<Option<TypedModelPatch>> {
91        if model.outlet_fact(node.inputs[0])?.datum_type == self.to {
92            return TypedModelPatch::shunt_one_op(model, node);
93        }
94        // linear_prec (fan-in=1, fan-out=1) rather than single_prec: swapping
95        // through a fan-out predecessor clones it, and the clone breaks
96        // downstream pattern detectors (e.g. Square+Reduce<Sum>+Mul fusion into
97        // Reduce<MeanOfSquares>, which then feeds RmsNorm detection).
98        //
99        // AxisOp is intentionally NOT in the predicate: pulling Cast above an
100        // AxisOp (Reshape/Move/Add/Rm) prevents the CUDA conversion from
101        // fusing the post-AxisOp Cast into the downstream GEMM-class kernel,
102        // leaving ~64 standalone CudaCast ops on OpenELM-270M (TG128 -4%).
103        if let Some(prec) = model.linear_prec(node.id)?
104            && (prec.op_is::<IntoShape>() || prec.op_is::<MultiBroadcastTo>())
105        {
106            let mut patch = TypedModelPatch::default();
107            let mut wire = tvec!(patch.tap_model(model, prec.inputs[0])?);
108            wire = patch.wire_node(&node.name, &node.op, &wire)?;
109            wire = patch.wire_node(&prec.name, &prec.op, &wire)?;
110            patch.shunt_outside(model, node.id.into(), wire[0])?;
111            return Ok(Some(patch));
112        }
113        Ok(None)
114    }
115
116    fn axes_mapping(
117        &self,
118        inputs: &[&TypedFact],
119        outputs: &[&TypedFact],
120    ) -> TractResult<AxesMapping> {
121        AxesMapping::natural(inputs, outputs)
122    }
123
124    fn change_axes(
125        &self,
126        model: &TypedModel,
127        node: &TypedNode,
128        _io: InOut,
129        change: &AxisOp,
130    ) -> TractResult<Option<AxisChangeConsequence>> {
131        Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
132    }
133
134    fn slice(
135        &self,
136        patch: &mut TypedModelPatch,
137        _model: &TypedModel,
138        node: &TypedNode,
139        _prefix: &str,
140        inputs: &[OutletId],
141        _output_axis: usize,
142        _start: &TDim,
143        _end: &TDim,
144    ) -> TractResult<Option<TVec<OutletId>>> {
145        patch.wire_node(&node.name, &node.op, inputs).map(Some)
146    }
147
148    as_op!();
149}
tract_core/ops/cast.rs

tract_core/ops/
cast.rs
