tract_gpu/ops/
pad.rs

1use crate::tensor::{DeviceTensorExt, IntoDevice};
2use tract_core::internal::*;
3use tract_core::ops::array::{Pad, PadMode};
4
5/// Constant padding via two `copy_nd`s: broadcast the pad value across the whole
6/// output, then drop the input into the interior. No dedicated kernel. Reflect/
7/// Edge modes are left on the host (see [`GpuPad::from_core`]).
8#[derive(Clone, Debug, PartialEq, Eq, Hash)]
9pub struct GpuPad {
10    pub pads: Vec<(usize, usize)>,
11    pub value: Arc<Tensor>,
12}
13
14impl GpuPad {
15    /// Build from a core `Pad`, or `None` when the mode isn't `Constant`.
16    pub fn from_core(op: &Pad) -> Option<Self> {
17        let PadMode::Constant(value) = &op.mode else { return None };
18        Some(Self { pads: op.pads.clone(), value: value.clone() })
19    }
20
21    fn output_shape<D: DimLike>(&self, input: &[D]) -> TVec<D> {
22        input.iter().zip(&self.pads).map(|(d, (a, b))| d.clone() + *a + *b).collect()
23    }
24}
25
26impl Op for GpuPad {
27    fn name(&self) -> StaticName {
28        "GpuPad".into()
29    }
30
31    op_as_typed_op!();
32}
33
34impl EvalOp for GpuPad {
35    op_out_of_plan!();
36
37    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
38        let input_value = args_1!(inputs);
39        let input = input_value.to_device_tensor()?;
40        let dt = input.datum_type();
41        let out_shape = self.output_shape(input.shape());
42
43        let output = crate::turn_handler::make_tensor_for_node(ctx, dt, &out_shape)?;
44
45        let ctx = crate::device::get_context()?;
46
47        // Fill the whole output with the pad value, broadcast from a scalar.
48        let value = self.value.cast_to_dt(dt)?.into_owned().into_device()?;
49        let zero_strides = vec![0isize; out_shape.len()];
50        ctx.copy_nd(&value, 0, &zero_strides, &output, 0, &out_shape, output.strides())?;
51
52        // Place the input at the interior offset.
53        if input.len() != 0 {
54            let interior: usize = self
55                .pads
56                .iter()
57                .enumerate()
58                .map(|(axis, (before, _))| before * output.strides()[axis] as usize)
59                .sum();
60            ctx.copy_nd(
61                input,
62                0,
63                input.strides(),
64                &output,
65                interior * dt.size_of(),
66                input.shape(),
67                output.strides(),
68            )?;
69        }
70        Ok(tvec![output.into_tensor().into_tvalue()])
71    }
72}
73
74impl TypedOp for GpuPad {
75    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
76        crate::utils::facts_to_device_facts(inputs, |facts| {
77            Ok(tvec!(facts[0].datum_type.fact(self.output_shape(&facts[0].shape.to_tvec()))))
78        })
79        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
80    }
81
82    as_op!();
83}
tract_gpu/ops/pad.rs

tract_gpu/ops/
pad.rs
