tract_gpu/ops/
concat.rs

1use crate::tensor::DeviceTensorExt;
2use tract_core::internal::*;
3
4#[derive(Clone, Debug, PartialEq, Eq, Hash)]
5pub struct GpuConcat {
6    pub axis: usize,
7}
8
9impl GpuConcat {
10    pub fn new(axis: usize) -> Self {
11        Self { axis }
12    }
13
14    pub fn offsets(&self, inputs: &[&TypedFact]) -> TractResult<Vec<TDim>> {
15        let mut offsets = vec![0.to_dim()];
16        for slice in inputs {
17            let len = slice.shape[self.axis].clone();
18            let offset = len + offsets.last().unwrap();
19            offsets.push(offset)
20        }
21        Ok(offsets)
22    }
23}
24
25impl Op for GpuConcat {
26    fn name(&self) -> StaticName {
27        "GpuConcat".into()
28    }
29
30    fn info(&self) -> TractResult<Vec<String>> {
31        Ok(vec![format!("axis: {}", self.axis)])
32    }
33
34    op_as_typed_op!();
35}
36
37impl EvalOp for GpuConcat {
38    op_out_of_plan!();
39
40    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
41        let inputs =
42            inputs.iter().map(|it| it.to_device_tensor()).collect::<TractResult<TVec<_>>>()?;
43
44        let mut output_shape = inputs[0].shape().to_vec();
45        output_shape[self.axis] = inputs.iter().map(|it| it.shape()[self.axis]).sum();
46        let output =
47            crate::turn_handler::make_tensor_for_node(ctx, inputs[0].datum_type(), &output_shape)?;
48
49        let ctx = crate::device::get_context()?;
50        let mut cursor = 0usize;
51        for input in &inputs {
52            let slice_len = input.shape()[self.axis];
53            if slice_len == 0 {
54                continue;
55            }
56            // Build zone shape (same as input shape for this slice)
57            let zone_shape = input.shape();
58            // Output offset along concat axis
59            let dst_offset =
60                cursor * output.strides()[self.axis] as usize * output.datum_type().size_of();
61
62            ctx.copy_nd(
63                input,
64                0,
65                input.strides(),
66                &output,
67                dst_offset,
68                zone_shape,
69                output.strides(),
70            )
71            .with_context(|| {
72                format!(
73                    "Error in concat dispatch for slice at offset {} (shape {:?})",
74                    cursor, zone_shape
75                )
76            })?;
77            cursor += slice_len;
78        }
79
80        Ok(tvec!(output.into_tensor().into_tvalue()))
81    }
82}
83
84impl TypedOp for GpuConcat {
85    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
86        crate::utils::facts_to_device_facts(inputs, |facts| {
87            let mut fact = facts[0].without_value();
88            for input in facts {
89                if input.rank() != fact.rank()
90                    || input
91                        .shape
92                        .iter()
93                        .zip(fact.shape.iter())
94                        .enumerate()
95                        .filter(|(ax, _)| *ax != self.axis)
96                        .any(|(_, (i, f))| i != f)
97                {
98                    bail!("Inconsistent {:?} inputs: {:?}", self, facts);
99                }
100            }
101            fact.shape.set(self.axis, self.offsets(facts)?.pop().unwrap());
102            Ok(tvec!(fact))
103        })
104        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
105    }
106
107    as_op!();
108}
tract_gpu/ops/concat.rs

tract_gpu/ops/
concat.rs
