tract_gpu/ops/
cast.rs

1use crate::tensor::DeviceTensorExt;
2use tract_core::internal::*;
3
4use crate::tensor::DeviceTensor;
5
6pub type DispatchCastFn = fn(&DeviceTensor, &DeviceTensor) -> TractResult<()>;
7
8#[derive(Clone)]
9pub struct GpuCast {
10    pub to: DatumType,
11    pub backend_name: &'static str,
12    pub dispatch: DispatchCastFn,
13    pub is_supported_dt: fn(DatumType) -> bool,
14}
15
16impl GpuCast {
17    pub fn new(
18        to: DatumType,
19        backend_name: &'static str,
20        dispatch: DispatchCastFn,
21        is_supported_dt: fn(DatumType) -> bool,
22    ) -> Option<Self> {
23        is_supported_dt(to).then_some(Self { to, backend_name, dispatch, is_supported_dt })
24    }
25
26    pub fn is_supported_dt(&self, dt: DatumType) -> bool {
27        (self.is_supported_dt)(dt)
28    }
29}
30
31impl std::fmt::Debug for GpuCast {
32    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
33        write!(f, "{}Cast({:?})", self.backend_name, self.to)
34    }
35}
36
37impl PartialEq for GpuCast {
38    fn eq(&self, other: &Self) -> bool {
39        self.backend_name == other.backend_name && self.to == other.to
40    }
41}
42
43impl Eq for GpuCast {}
44
45impl std::hash::Hash for GpuCast {
46    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
47        self.backend_name.hash(state);
48        self.to.hash(state);
49    }
50}
51
52impl Op for GpuCast {
53    fn name(&self) -> StaticName {
54        format!("{}Cast", self.backend_name).into()
55    }
56
57    op_as_typed_op!();
58}
59
60impl EvalOp for GpuCast {
61    op_out_of_plan!();
62
63    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
64        let input_value = args_1!(inputs);
65        let input = input_value.to_device_tensor()?;
66        if input.datum_type() == self.to {
67            Ok(tvec!(input_value))
68        } else {
69            let output = crate::turn_handler::make_tensor_for_node(ctx, self.to, input.shape())?;
70            (self.dispatch)(input, &output)?;
71            Ok(tvec![output.into_tensor().into_tvalue()])
72        }
73    }
74}
75
76impl TypedOp for GpuCast {
77    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
78        crate::utils::facts_to_device_facts(inputs, |facts| {
79            Ok(tvec!(self.to.fact(facts[0].shape.clone())))
80        })
81        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
82    }
83
84    as_op!();
85}
tract_gpu/ops/cast.rs

tract_gpu/ops/
cast.rs
