tract_gpu/ops/
iff.rs

1use crate::tensor::{DeviceTensor, DeviceTensorExt};
2use derive_new::new;
3use tract_core::broadcast::multi_broadcast;
4use tract_core::internal::*;
5
6static IFF_MAX_RANK: usize = 5;
7
8/// Dispatch function for the iff (select) kernel.
9/// Args: cond, then, else tensors with pre-computed broadcast strides,
10/// output tensor, output shape and strides. All strides are padded to IFF_MAX_RANK.
11pub type DispatchIffFn = fn(
12    cond: &DeviceTensor,
13    then_value: &DeviceTensor,
14    else_value: &DeviceTensor,
15    cond_strides: &[isize],
16    then_strides: &[isize],
17    else_strides: &[isize],
18    output: &DeviceTensor,
19    output_shape: &[usize],
20    output_strides: &[isize],
21) -> TractResult<()>;
22
23#[derive(Clone, new)]
24pub struct GpuIff {
25    pub backend_name: &'static str,
26    pub dispatch: DispatchIffFn,
27}
28
29impl std::fmt::Debug for GpuIff {
30    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
31        write!(f, "{}Iff", self.backend_name)
32    }
33}
34
35impl PartialEq for GpuIff {
36    fn eq(&self, other: &Self) -> bool {
37        self.backend_name == other.backend_name
38    }
39}
40
41impl Eq for GpuIff {}
42
43impl std::hash::Hash for GpuIff {
44    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
45        self.backend_name.hash(state);
46    }
47}
48
49impl Op for GpuIff {
50    fn name(&self) -> StaticName {
51        format!("{}Iff", self.backend_name).into()
52    }
53
54    op_as_typed_op!();
55}
56
57impl EvalOp for GpuIff {
58    op_out_of_plan!();
59
60    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
61        let (cond_val, then_val, else_val) = args_3!(inputs);
62
63        let cond = cond_val.to_device_tensor()?;
64        let then_t = then_val.to_device_tensor()?;
65        let else_t = else_val.to_device_tensor()?;
66        ensure!(cond.rank() == then_t.rank());
67        ensure!(cond.rank() == else_t.rank());
68        ensure!(then_t.datum_type() == else_t.datum_type());
69
70        let out_shape = multi_broadcast(&[cond.shape(), then_t.shape(), else_t.shape()])
71            .context("No broadcasting solution found")?;
72        let out_dt = then_t.datum_type();
73        let output = crate::turn_handler::make_tensor_for_node(ctx, out_dt, &out_shape)?;
74
75        if output.len() > 0 {
76            let rank = cond.rank();
77            ensure!(rank <= IFF_MAX_RANK);
78            let rank_pad = IFF_MAX_RANK - rank;
79
80            let mut padded_cond_strides = [0isize; IFF_MAX_RANK];
81            let mut padded_then_strides = [0isize; IFF_MAX_RANK];
82            let mut padded_else_strides = [0isize; IFF_MAX_RANK];
83            let mut padded_out_shape = [1usize; IFF_MAX_RANK];
84            let mut padded_out_strides = [0isize; IFF_MAX_RANK];
85
86            for axis in 0..rank {
87                padded_out_shape[rank_pad + axis] = output.shape()[axis];
88                padded_out_strides[rank_pad + axis] = output.strides()[axis];
89                padded_cond_strides[rank_pad + axis] = if cond.shape()[axis] < output.shape()[axis]
90                {
91                    0
92                } else {
93                    cond.strides()[axis]
94                };
95                padded_then_strides[rank_pad + axis] =
96                    if then_t.shape()[axis] < output.shape()[axis] {
97                        0
98                    } else {
99                        then_t.strides()[axis]
100                    };
101                padded_else_strides[rank_pad + axis] =
102                    if else_t.shape()[axis] < output.shape()[axis] {
103                        0
104                    } else {
105                        else_t.strides()[axis]
106                    };
107            }
108
109            (self.dispatch)(
110                cond,
111                then_t,
112                else_t,
113                &padded_cond_strides,
114                &padded_then_strides,
115                &padded_else_strides,
116                &output,
117                &padded_out_shape,
118                &padded_out_strides,
119            )
120            .with_context(|| "Error while dispatching eval for Iff")?;
121        }
122        Ok(tvec!(output.into_tensor().into_tvalue()))
123    }
124}
125
126impl TypedOp for GpuIff {
127    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
128        crate::utils::facts_to_device_facts(inputs, |inputs| {
129            let out_shape =
130                multi_broadcast(&[&*inputs[0].shape, &*inputs[1].shape, &*inputs[2].shape])
131                    .context("No broadcasting solution found")?;
132            let out_dt = inputs[1].datum_type;
133            Ok(tvec!(out_dt.fact(out_shape)))
134        })
135    }
136
137    as_op!();
138}
tract_gpu/ops/iff.rs

tract_gpu/ops/
iff.rs
