tract_gpu/
sync.rs

1use crate::fact::{DeviceFact, DeviceTypedFactExt};
2use crate::tensor::{DeviceTensor, DeviceTensorExt, IntoDevice, LazyHostStorage};
3use derive_new::new;
4use std::collections::HashMap;
5use std::fmt;
6use std::sync::Arc;
7use tract_core::internal::*;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum DeviceSyncKind {
11    ToHost,
12    ToDevice,
13}
14
15impl fmt::Display for DeviceSyncKind {
16    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
17        write!(f, "{self:?}")
18    }
19}
20
21#[derive(Debug, Clone, new, Copy, PartialEq, Eq, Hash)]
22pub struct DeviceSync {
23    pub kind: DeviceSyncKind,
24}
25
26impl Op for DeviceSync {
27    fn name(&self) -> StaticName {
28        format!("DeviceSync{}", self.kind).into()
29    }
30
31    op_as_typed_op!();
32}
33
34impl EvalOp for DeviceSync {
35    op_out_of_plan!();
36
37    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
38        let input = args_1!(inputs);
39        match self.kind {
40            DeviceSyncKind::ToHost => {
41                let device_tensor = input.to_device_tensor()?;
42                match device_tensor {
43                    // An owned device tensor of plain layout crosses the
44                    // boundary lazily: the result is a host tensor, but the
45                    // bytes only come back if someone reads them. A caller that
46                    // hands it straight to the next run never pays the
47                    // transfer.
48                    DeviceTensor::Owned(_) if !device_tensor.is_exotic() => Ok(tvec![
49                        LazyHostStorage::new(device_tensor.clone())?.into_tensor().into_tvalue()
50                    ]),
51                    // A view borrows a buffer whose contents are only good for
52                    // the turn, so it cannot outlive it; an exotic tensor comes back as its own
53                    // storage -- block-quant weights come back block-quant --
54                    // which a lazy host tensor has no way to stand in for.
55                    // Both copy out now.
56                    _ => {
57                        let tensor = device_tensor
58                            .to_host()
59                            .with_context(|| "Error while syncing device tensor to host")?;
60                        Ok(tvec![tensor.into_tvalue()])
61                    }
62                }
63            }
64            DeviceSyncKind::ToDevice => {
65                // A tensor that came back from a previous run still on device
66                // goes straight back in, with no round trip either way.
67                if let Some(device) =
68                    input.storage_as::<LazyHostStorage>().and_then(|lazy| lazy.device())
69                {
70                    return Ok(tvec![device.clone().into_tensor().into()]);
71                }
72                let device_input = if let Some(t) = input.as_arc_tensor() {
73                    Arc::clone(t).into_device()?
74                } else {
75                    input.into_tensor().into_device()?
76                };
77                Ok(tvec![device_input.into_tensor().into()])
78            }
79        }
80    }
81}
82
83impl TypedOp for DeviceSync {
84    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
85        let input = inputs[0];
86        match self.kind {
87            DeviceSyncKind::ToHost => {
88                let mut typed_fact = input
89                    .to_device_fact()
90                    .with_context(|| {
91                        "Cannot sync to Host a tensor without DeviceFact as metadata in its TypedFact"
92                    })?
93                    .clone()
94                    .into_typed_fact();
95                if let Some(konst) = input.konst.clone() {
96                    if let Some(dt) = konst.as_device_tensor() {
97                        typed_fact.konst = Some(dt.to_host()?);
98                    } else {
99                        typed_fact.konst = Some(konst);
100                    }
101                }
102                Ok(tvec!(typed_fact))
103            }
104            DeviceSyncKind::ToDevice => {
105                ensure!(
106                    input.as_device_fact().is_none(),
107                    "Cannot sync to Device a tensor already on Device"
108                );
109                Ok(tvec![DeviceFact::from_host(input.clone())?.into_exotic_fact()])
110            }
111        }
112    }
113
114    as_op!();
115}
116
117/// Map node inputs through the translation mapping, inserting DeviceSync nodes
118/// where needed to move tensors to/from the device.
119pub fn sync_inputs_if_required(
120    model: &mut TypedModel,
121    node: &TypedNode,
122    mapping: &HashMap<OutletId, OutletId>,
123    sync_kind: DeviceSyncKind,
124) -> TractResult<TVec<OutletId>> {
125    let mut mapped_inputs = tvec![];
126    for (i_idx, i) in node.inputs.iter().enumerate() {
127        let in_fact = model.outlet_fact_mut(mapping[i])?;
128        match sync_kind {
129            DeviceSyncKind::ToHost if in_fact.as_device_fact().is_some() => {
130                mapped_inputs.push(
131                    model.wire_node(
132                        format!("{}.to-cpu-{i_idx}", node.name),
133                        DeviceSync::new(sync_kind),
134                        &[mapping[i]],
135                    )?[0],
136                );
137            }
138            DeviceSyncKind::ToDevice if in_fact.as_device_fact().is_none() => {
139                if let Some(ref konst) = in_fact.konst
140                    && konst.as_device_tensor().is_none()
141                {
142                    let device_konst = konst.as_ref().clone().into_device()?.into_tensor();
143                    let device_fact = DeviceFact::from_host(in_fact.clone())?;
144
145                    *in_fact = device_fact.into_exotic_fact();
146
147                    in_fact.konst = Some(Arc::new(device_konst));
148                    mapped_inputs.push(mapping[i]);
149                    continue;
150                }
151                ensure!(
152                    in_fact.datum_type.is_copy(),
153                    "Only copy DatumType can be sync to Device: {:?}",
154                    in_fact.datum_type
155                );
156
157                mapped_inputs.push(
158                    model.wire_node(
159                        format!("{}.to-device-{i_idx}", node.name),
160                        DeviceSync::new(sync_kind),
161                        &[mapping[i]],
162                    )?[0],
163                );
164            }
165            _ => mapped_inputs.push(mapping[i]),
166        }
167    }
168    Ok(mapped_inputs)
169}
170
171/// For model outputs that are on device, insert DeviceSync nodes to move them back to host.
172pub fn sync_model_outputs_if_required(
173    src: &TypedModel,
174    node: &TypedNode,
175    target: &mut TypedModel,
176    target_node_outlet_ids: TVec<OutletId>,
177) -> TractResult<TVec<OutletId>> {
178    let mut outputs = tvec![];
179    for (o_idx, o) in target_node_outlet_ids.into_iter().enumerate() {
180        let is_src_output = src.outputs.contains(&OutletId::new(node.id, o_idx));
181        if target.outlet_fact(o)?.as_device_fact().is_some() && is_src_output {
182            let sync_output = target.wire_node(
183                format!("{}.to-host-{o_idx}-out", node.name),
184                DeviceSync::new(DeviceSyncKind::ToHost),
185                &[o],
186            )?[0];
187            outputs.push(sync_output);
188        } else {
189            outputs.push(o)
190        }
191    }
192    Ok(outputs)
193}
tract_gpu/sync.rs

tract_gpu/
sync.rs
