tract_gpu/
sync.rs

1use crate::fact::{DeviceFact, DeviceTypedFactExt};
2use crate::tensor::{DeviceTensorExt, IntoDevice};
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
43                let tensor = device_tensor
44                    .to_host()
45                    .with_context(|| "Error while syncing device tensor to host")?;
46                Ok(tvec![tensor.into_tvalue()])
47            }
48            DeviceSyncKind::ToDevice => {
49                let device_input = if let Some(t) = input.as_arc_tensor() {
50                    Arc::clone(t).into_device()?
51                } else {
52                    input.into_tensor().into_device()?
53                };
54                Ok(tvec![device_input.into_tensor().into()])
55            }
56        }
57    }
58}
59
60impl TypedOp for DeviceSync {
61    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
62        let input = inputs[0];
63        match self.kind {
64            DeviceSyncKind::ToHost => {
65                let mut typed_fact = input
66                    .to_device_fact()
67                    .with_context(|| {
68                        "Cannot sync to Host a tensor without DeviceFact as metadata in its TypedFact"
69                    })?
70                    .clone()
71                    .into_typed_fact();
72                if let Some(konst) = input.konst.clone() {
73                    if let Some(dt) = konst.as_device_tensor() {
74                        typed_fact.konst = Some(dt.to_host()?);
75                    } else {
76                        typed_fact.konst = Some(konst);
77                    }
78                }
79                Ok(tvec!(typed_fact))
80            }
81            DeviceSyncKind::ToDevice => {
82                ensure!(
83                    input.as_device_fact().is_none(),
84                    "Cannot sync to Device a tensor already on Device"
85                );
86                Ok(tvec![DeviceFact::from_host(input.clone())?.into_exotic_fact()])
87            }
88        }
89    }
90
91    as_op!();
92}
93
94/// Map node inputs through the translation mapping, inserting DeviceSync nodes
95/// where needed to move tensors to/from the device.
96pub fn sync_inputs_if_required(
97    model: &mut TypedModel,
98    node: &TypedNode,
99    mapping: &HashMap<OutletId, OutletId>,
100    sync_kind: DeviceSyncKind,
101) -> TractResult<TVec<OutletId>> {
102    let mut mapped_inputs = tvec![];
103    for (i_idx, i) in node.inputs.iter().enumerate() {
104        let in_fact = model.outlet_fact_mut(mapping[i])?;
105        match sync_kind {
106            DeviceSyncKind::ToHost if in_fact.as_device_fact().is_some() => {
107                mapped_inputs.push(
108                    model.wire_node(
109                        format!("{}.to-cpu-{i_idx}", node.name),
110                        DeviceSync::new(sync_kind),
111                        &[mapping[i]],
112                    )?[0],
113                );
114            }
115            DeviceSyncKind::ToDevice if in_fact.as_device_fact().is_none() => {
116                if let Some(ref konst) = in_fact.konst
117                    && konst.as_device_tensor().is_none()
118                {
119                    let device_konst = konst.as_ref().clone().into_device()?.into_tensor();
120                    let device_fact = DeviceFact::from_host(in_fact.clone())?;
121
122                    *in_fact = device_fact.into_exotic_fact();
123
124                    in_fact.konst = Some(Arc::new(device_konst));
125                    mapped_inputs.push(mapping[i]);
126                    continue;
127                }
128                ensure!(
129                    in_fact.datum_type.is_copy(),
130                    "Only copy DatumType can be sync to Device: {:?}",
131                    in_fact.datum_type
132                );
133
134                mapped_inputs.push(
135                    model.wire_node(
136                        format!("{}.to-device-{i_idx}", node.name),
137                        DeviceSync::new(sync_kind),
138                        &[mapping[i]],
139                    )?[0],
140                );
141            }
142            _ => mapped_inputs.push(mapping[i]),
143        }
144    }
145    Ok(mapped_inputs)
146}
147
148/// For model outputs that are on device, insert DeviceSync nodes to move them back to host.
149pub fn sync_model_outputs_if_required(
150    src: &TypedModel,
151    node: &TypedNode,
152    target: &mut TypedModel,
153    target_node_outlet_ids: TVec<OutletId>,
154) -> TractResult<TVec<OutletId>> {
155    let mut outputs = tvec![];
156    for (o_idx, o) in target_node_outlet_ids.into_iter().enumerate() {
157        let is_src_output = src.outputs.contains(&OutletId::new(node.id, o_idx));
158        if target.outlet_fact(o)?.as_device_fact().is_some() && is_src_output {
159            let sync_output = target.wire_node(
160                format!("{}.to-host-{o_idx}-out", node.name),
161                DeviceSync::new(DeviceSyncKind::ToHost),
162                &[o],
163            )?[0];
164            outputs.push(sync_output);
165        } else {
166            outputs.push(o)
167        }
168    }
169    Ok(outputs)
170}
tract_gpu/sync.rs

tract_gpu/
sync.rs
