tract_gpu/ops/
gdn_recurrent.rs

1use crate::tensor::{DeviceTensor, DeviceTensorExt};
2use crate::turn_handler::make_tensor_for_node;
3use tract_core::internal::*;
4
5pub type DispatchGdnRecurrentFn = fn(
6    &DeviceTensor,
7    &DeviceTensor,
8    &DeviceTensor,
9    &DeviceTensor,
10    &DeviceTensor,
11    &DeviceTensor,
12    &DeviceTensor,
13    &DeviceTensor,
14) -> TractResult<()>;
15
16#[derive(Clone, Debug)]
17pub struct GpuGatedDeltaNetRecurrent {
18    pub backend_name: &'static str,
19    pub dispatch: DispatchGdnRecurrentFn,
20}
21
22impl PartialEq for GpuGatedDeltaNetRecurrent {
23    fn eq(&self, other: &Self) -> bool {
24        self.backend_name == other.backend_name
25    }
26}
27impl Eq for GpuGatedDeltaNetRecurrent {}
28impl std::hash::Hash for GpuGatedDeltaNetRecurrent {
29    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
30        self.backend_name.hash(state);
31    }
32}
33
34impl Op for GpuGatedDeltaNetRecurrent {
35    fn name(&self) -> StaticName {
36        format!("{}GatedDeltaNetRecurrent", self.backend_name).into()
37    }
38    op_as_typed_op!();
39}
40
41impl EvalOp for GpuGatedDeltaNetRecurrent {
42    op_out_of_plan!();
43
44    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
45        ensure!(inputs.len() == 6);
46        let tensors = inputs
47            .iter()
48            .map(|value| value.to_device_tensor())
49            .collect::<TractResult<TVec<_>>>()?;
50        let output = make_tensor_for_node(ctx, DatumType::F16, tensors[0].shape())?;
51        // The memory arena is keyed by node, so a second output cannot use the
52        // same arena slot as the first one.
53        let final_state = DeviceTensor::uninitialized_dt(DatumType::F32, tensors[5].shape())?;
54        (self.dispatch)(
55            tensors[0],
56            tensors[1],
57            tensors[2],
58            tensors[3],
59            tensors[4],
60            tensors[5],
61            &output,
62            &final_state,
63        )?;
64        Ok(tvec![output.into_tensor().into_tvalue(), final_state.into_tensor().into_tvalue()])
65    }
66}
67
68impl TypedOp for GpuGatedDeltaNetRecurrent {
69    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
70        crate::utils::facts_to_device_facts(inputs, |facts| {
71            ensure!(facts.len() == 6);
72            ensure!(facts[0].datum_type == DatumType::F16);
73            ensure!(facts[1].datum_type == DatumType::F16);
74            ensure!(facts[2].datum_type == DatumType::F16);
75            ensure!(facts[3].datum_type == DatumType::F32);
76            ensure!(facts[4].datum_type == DatumType::F16);
77            ensure!(facts[5].datum_type == DatumType::F32);
78            Ok(tvec![facts[0].without_value(), facts[5].without_value()])
79        })
80        .with_context(|| format!("invalid facts for {}", self.name()))
81    }
82    as_op!();
83}
tract_gpu/ops/gdn_recurrent.rs

tract_gpu/ops/
gdn_recurrent.rs
