tract_gpu/ops/
broadcast.rs

1use crate::tensor::DeviceTensorExt;
2use crate::utils::compute_broadcast_strides;
3use tract_core::internal::*;
4
5#[derive(Clone, Debug, PartialEq, Eq, Hash)]
6pub struct GpuMultiBroadcastTo {
7    pub shape: ShapeFact,
8}
9
10impl GpuMultiBroadcastTo {
11    pub fn new(shape: ShapeFact) -> Self {
12        Self { shape }
13    }
14}
15
16impl Op for GpuMultiBroadcastTo {
17    fn name(&self) -> StaticName {
18        "GpuMultiBroadcastTo".into()
19    }
20
21    op_as_typed_op!();
22}
23
24impl EvalOp for GpuMultiBroadcastTo {
25    op_out_of_plan!();
26
27    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
28        let input_value = args_1!(inputs);
29        let input = input_value.to_device_tensor()?;
30        let shape = self.shape.eval_to_usize(ctx.symbols)?;
31        let output = crate::turn_handler::make_tensor_for_node(ctx, input.datum_type(), &shape)?;
32
33        // Pad input shape/strides to output rank for broadcasting.  The padded
34        // axes have dim 1, so `compute_broadcast_strides` zeroes their stride
35        // whatever it is given -- which is what lets a rank-0 input, whose
36        // `strides()` is empty, broadcast at all.
37        let pad_stride = input.strides().first().copied().unwrap_or(1);
38        let mut input_strides = vec![pad_stride; output.rank() - input.rank()];
39        input_strides.extend(input.strides());
40        let mut input_shape = vec![1usize; output.rank() - input.rank()];
41        input_shape.extend(input.shape());
42        let broadcast_strides: TVec<isize> =
43            compute_broadcast_strides(&input_shape, &input_strides)?;
44
45        let ctx = crate::device::get_context()?;
46        ctx.copy_nd(input, 0, &broadcast_strides, &output, 0, output.shape(), output.strides())?;
47        Ok(tvec![output.into_tensor().into_tvalue()])
48    }
49}
50
51impl TypedOp for GpuMultiBroadcastTo {
52    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
53        crate::utils::facts_to_device_facts(inputs, |facts| {
54            let mut fact = facts[0].datum_type.fact(self.shape.clone());
55            fact.uniform.clone_from(&inputs[0].uniform);
56            Ok(tvec!(fact))
57        })
58        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
59    }
60
61    as_op!();
62}
tract_gpu/ops/broadcast.rs

tract_gpu/ops/
broadcast.rs
