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_ram()?;
52            let mut tmp_plain = tmp.try_as_plain_ram_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 fact.uniform_tdim.is_none() {
71            // a symbolic scalar konst or uniform keeps its value as metadata:
72            // the concrete cast cannot evaluate it, but downstream consumers
73            // (a Range limit, say) can still reason about it
74            for t in
75                [inputs[0].uniform.as_deref(), inputs[0].konst.as_deref()].into_iter().flatten()
76            {
77                if t.datum_type() == TDim::datum_type() && t.volume() == 1 {
78                    fact.uniform_tdim = Some(
79                        t.try_as_plain_ram()?.to_scalar::<TDim>()?.eval(&SymbolValues::default()),
80                    );
81                    break;
82                }
83            }
84        }
85        if let Some(u) = &inputs[0].uniform
86            && let Ok(cast_u) = u.cast_to_dt(self.to)
87        {
88            fact.uniform = Some(std::sync::Arc::new(cast_u.into_owned()));
89        }
90        Ok(tvec!(fact))
91    }
92
93    fn input_roi(
94        &self,
95        model: &TypedModel,
96        node: &TypedNode,
97    ) -> TractResult<Option<TVec<Option<TDim>>>> {
98        crate::optim::propagate_roi::bubble_roi(model, node)
99    }
100
101    fn declutter(
102        &self,
103        model: &TypedModel,
104        node: &TypedNode,
105    ) -> TractResult<Option<TypedModelPatch>> {
106        if model.outlet_fact(node.inputs[0])?.datum_type == self.to {
107            return TypedModelPatch::shunt_one_op(model, node);
108        }
109        // linear_prec (fan-in=1, fan-out=1) rather than single_prec: swapping
110        // through a fan-out predecessor clones it, and the clone breaks
111        // downstream pattern detectors (e.g. Square+Reduce<Sum>+Mul fusion into
112        // Reduce<MeanOfSquares>, which then feeds RmsNorm detection).
113        //
114        // AxisOp is intentionally NOT in the predicate: pulling Cast above an
115        // AxisOp (Reshape/Move/Add/Rm) prevents the CUDA conversion from
116        // fusing the post-AxisOp Cast into the downstream GEMM-class kernel,
117        // leaving ~64 standalone CudaCast ops on OpenELM-270M (TG128 -4%).
118        if let Some(prec) = model.linear_prec(node.id)?
119            && (prec.op_is::<IntoShape>() || prec.op_is::<MultiBroadcastTo>())
120        {
121            let mut patch = TypedModelPatch::default();
122            let mut wire = tvec!(patch.tap_model(model, prec.inputs[0])?);
123            wire = patch.wire_node(&node.name, &node.op, &wire)?;
124            wire = patch.wire_node(&prec.name, &prec.op, &wire)?;
125            patch.shunt_outside(model, node.id.into(), wire[0])?;
126            return Ok(Some(patch));
127        }
128        Ok(None)
129    }
130
131    fn axes_mapping(
132        &self,
133        inputs: &[&TypedFact],
134        outputs: &[&TypedFact],
135    ) -> TractResult<AxesMapping> {
136        AxesMapping::natural(inputs, outputs)
137    }
138
139    fn change_axes(
140        &self,
141        model: &TypedModel,
142        node: &TypedNode,
143        _io: InOut,
144        change: &AxisOp,
145    ) -> TractResult<Option<AxisChangeConsequence>> {
146        Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
147    }
148
149    fn slice(
150        &self,
151        patch: &mut TypedModelPatch,
152        _model: &TypedModel,
153        node: &TypedNode,
154        _prefix: &str,
155        inputs: &[OutletId],
156        _output_axis: usize,
157        _start: &TDim,
158        _end: &TDim,
159    ) -> TractResult<Option<TVec<OutletId>>> {
160        patch.wire_node(&node.name, &node.op, inputs).map(Some)
161    }
162
163    as_op!();
164}
tract_core/ops/cast.rs

tract_core/ops/
cast.rs
