tract_gpu/ops/
rotate_half.rs

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

tract_gpu/ops/
rotate_half.rs
