tract_gpu/ops/
slice.rs

1use crate::tensor::DeviceTensorExt;
2use crate::utils::compute_broadcast_strides;
3use tract_core::internal::*;
4use tract_core::ops::array::Slice;
5
6#[derive(Clone, Debug, PartialEq, Eq, Hash)]
7pub struct GpuSlice {
8    pub inner: Slice,
9}
10
11impl GpuSlice {
12    pub fn new(inner: Slice) -> Self {
13        Self { inner }
14    }
15}
16
17impl Op for GpuSlice {
18    fn name(&self) -> StaticName {
19        "GpuSlice".into()
20    }
21
22    fn info(&self) -> TractResult<Vec<String>> {
23        self.inner.info()
24    }
25
26    op_as_typed_op!();
27}
28
29impl EvalOp for GpuSlice {
30    op_out_of_plan!();
31
32    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
33        let input_value = args_1!(inputs);
34        let input = input_value.to_device_tensor()?;
35
36        let start = self.inner.start.eval(ctx.symbols).to_usize()?;
37        let end = self.inner.end.eval(ctx.symbols).to_usize()?;
38        let axis = self.inner.axis;
39
40        let input_shape = input.shape();
41        let input_strides = input.strides();
42        let input_dt = input.datum_type();
43
44        ensure!(
45            end <= input_shape[axis] && start <= end,
46            "Invalid range {}..{} for slicing {:?} on axis {}",
47            start,
48            end,
49            input,
50            axis
51        );
52
53        let mut o_shape: TVec<usize> = input_shape.into();
54        o_shape[axis] = end - start;
55
56        let offset = (start * input_strides[axis] as usize) * input_dt.size_of();
57
58        let output = crate::turn_handler::make_tensor_for_node(ctx, input.datum_type(), &o_shape)?;
59
60        if o_shape[axis] != 0 {
61            // Slice uses same strides as input (broadcast strides with matching shapes)
62            let broadcast_strides: TVec<isize> =
63                compute_broadcast_strides(&o_shape, input_strides)?;
64            let ctx = crate::device::get_context()?;
65            ctx.copy_nd(
66                input,
67                offset,
68                &broadcast_strides,
69                &output,
70                0,
71                output.shape(),
72                output.strides(),
73            )?;
74        }
75        Ok(tvec![output.into_tensor().into_tvalue()])
76    }
77}
78
79impl TypedOp for GpuSlice {
80    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
81        crate::utils::facts_to_device_facts(inputs, |facts| self.inner.output_facts(facts))
82            .with_context(|| format!("Error while computing facts for {:?}", self.name()))
83    }
84
85    fn set_symbols(
86        &self,
87        _source: &TypedModel,
88        node: &TypedNode,
89        target: &mut TypedModel,
90        mapping: &HashMap<OutletId, OutletId>,
91        subs: &HashMap<Symbol, TDim>,
92    ) -> TractResult<TVec<OutletId>> {
93        let op = GpuSlice {
94            inner: Slice {
95                axis: self.inner.axis,
96                start: self.inner.start.substitute_all(subs)?,
97                end: self.inner.end.substitute_all(subs)?,
98            },
99        };
100        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
101        target.wire_node(&node.name, op, &inputs)
102    }
103
104    as_op!();
105}
tract_gpu/ops/slice.rs

tract_gpu/ops/
slice.rs
