tract_core/optim/
change_axes.rs

1use super::OptimizerSession;
2use super::TypedPass;
3use crate::internal::*;
4use crate::model::*;
5use crate::ops::dummy::Dummy;
6use crate::ops::einsum::EinSum;
7use crate::ops::konst::Const;
8use std::collections::hash_map::Entry;
9use std::collections::HashSet;
10use std::fmt::Debug;
11
12use crate::ops::change_axes::*;
13
14#[derive(Clone, Default)]
15pub struct ChangeAxes(HashSet<crate::ops::change_axes::AxisChange>, usize);
16
17impl Debug for ChangeAxes {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        write!(f, "ChangeAxes")
20    }
21}
22
23impl TypedPass for ChangeAxes {
24    fn reset(&mut self) -> TractResult<()> {
25        self.0.clear();
26        self.1 = 0;
27        Ok(())
28    }
29    fn next(
30        &mut self,
31        _session: &mut OptimizerSession,
32        model: &TypedModel,
33    ) -> TractResult<Option<TypedModelPatch>> {
34        let mut explored: HashSet<AxisChange> = Default::default();
35        let mut interfaces = model.output_outlets()?.to_vec();
36        interfaces.extend(model.input_outlets()?.iter());
37        for node in &model.nodes[self.1..] {
38            if node.op_is::<Dummy>() {
39                continue;
40            }
41            for suggestion in node.op.suggested_axis_changes()? {
42                let outlet = suggestion.0.as_outlet(node);
43                let change = AxisChange { outlet, op: suggestion.1 };
44                if self.0.insert(change.clone()) {
45                    if let Some((patch, _)) =
46                        change_axes(model, &change, &interfaces, &[], &mut explored).with_context(
47                            || format!("Making patch for {change:?} from {node}"),
48                        )?
49                    {
50                        self.1 = node.id;
51                        return Ok(Some(patch));
52                    }
53                }
54            }
55            for (slot, fact) in node.outputs.iter().enumerate() {
56                for (ix, dim) in fact.fact.shape.iter().enumerate() {
57                    if dim.is_one() {
58                        let change =
59                            AxisChange { outlet: OutletId::new(node.id, slot), op: AxisOp::Rm(ix) };
60                        if self.0.insert(change.clone()) {
61                            if let Some((patch, _)) =
62                                change_axes(model, &change, &interfaces, &[], &mut explored)
63                                    .with_context(|| {
64                                        format!("Making patch for {change:?} from {node}")
65                                    })?
66                            {
67                                self.1 = node.id;
68                                return Ok(Some(patch));
69                            }
70                        }
71                    }
72                }
73            }
74        }
75        Ok(None)
76    }
77}
78
79#[allow(clippy::type_complexity)]
80pub fn change_axes(
81    model: &TypedModel,
82    change: &AxisChange,
83    locked: &[OutletId],
84    bounds: &[TVec<OutletId>],
85    explored: &mut HashSet<AxisChange>,
86) -> TractResult<Option<(TypedModelPatch, TVec<(InOut, AxisOp)>)>> {
87    if explored.contains(change) {
88        debug!("  Not considering change because deja vu {change:?}");
89        return Ok(None);
90    }
91    if model
92        .node(change.outlet.node)
93        .op_as::<Const>()
94        .is_some_and(|c| c.val().volume() == 1 && c.val().datum_type() != Opaque::datum_type())
95    {
96        debug!("  Not considering change from const {change:?}");
97        return Ok(None);
98    }
99    debug!("  Considering change {change:?}");
100    let mut todo_changes = vec![(change.clone(), None)];
101    let mut changed_wires: HashMap<TVec<OutletId>, AxisOp> = HashMap::new();
102    let bound_outlets = |o: OutletId| -> TVec<OutletId> {
103        bounds.iter().find(|b| b.contains(&o)).cloned().unwrap_or_else(|| tvec!(o))
104    };
105    changed_wires.insert(bound_outlets(change.outlet), change.op.clone());
106    let mut changed_ops: HashMap<usize, Box<dyn TypedOp>> = HashMap::new();
107    let mut rewired_scalar_input: HashMap<InletId, (OutletId, AxisOp)> = Default::default();
108    while let Some((change, emitter)) = todo_changes.pop() {
109        if !explored.insert(change.clone()) {
110            return Ok(None);
111        }
112        let outlet_group = bound_outlets(change.outlet);
113        for &outlet in &outlet_group {
114            if locked.contains(&outlet) {
115                debug!("    Change {change:?} blocked by locked interface {outlet:?}");
116                return Ok(None);
117            }
118            let mut interfaces: Vec<(usize, InOut)> = vec![(outlet.node, InOut::Out(outlet.slot))];
119            for inlet in model.outlet_successors(outlet) {
120                interfaces.push((inlet.node, InOut::In(inlet.slot)));
121            }
122            for (node_id, io) in interfaces {
123                if Some(node_id) == emitter {
124                    continue;
125                }
126                let node = model.node(node_id);
127                // if this is a revisit...
128                let op = if let Some(op) = changed_ops.get(&node_id) {
129                    trace!("  Change {:?} revisiting {}", change, model.node(node_id));
130                    if op.is::<EinSum>() {
131                        // FIXME Einsum can swallow any combination of axis change on all interfaces
132                        op
133                    } else {
134                        debug!("  Change {:?} blocked: revisiting {}", change, model.node(node_id));
135                        return Ok(None);
136                    }
137                } else {
138                    &node.op
139                };
140                let more = op
141                    .change_axes(model, node, io, &change.op)
142                    .with_context(|| format!("Propagating {change:?} to node {node}"))?;
143                if more.is_none() {
144                    debug!("    Propagation of {change:?} blocked by {node}");
145                    return Ok(None);
146                }
147                let AxisChangeConsequence { substitute_op, wire_changes } = more.unwrap();
148                trace!("    Change {:?} enters {} from {:?}", change.op, node, io);
149                trace!("       propagates as {wire_changes:?}");
150                if let Some(op) = substitute_op {
151                    trace!("       replace op by {op:?}");
152                    changed_ops.insert(node.id, op);
153                }
154                for (wire, op) in wire_changes.into_iter() {
155                    let outlet = wire.as_outlet(node);
156                    // stop upstram propagation to a scalar constant: we will clone it and alter it
157                    // at patch generation time
158                    if let InOut::In(inlet) = wire {
159                        if model
160                            .node(outlet.node)
161                            .op_as::<Const>()
162                            .is_some_and(|k| k.val().volume() == 1)
163                        {
164                            rewired_scalar_input.insert(InletId::new(node.id, inlet), (outlet, op));
165                            continue;
166                        }
167                    }
168                    let outlet_group = bound_outlets(wire.as_outlet(node));
169                    match changed_wires.entry(outlet_group.clone()) {
170                        Entry::Vacant(entry) => {
171                            trace!(
172                                "         {wire:?} {op:?} change on {outlet_group:?} is new"
173                            );
174                            entry.insert(op.clone());
175                            todo_changes
176                                .push((AxisChange { outlet: outlet_group[0], op }, Some(node_id)));
177                        }
178                        Entry::Occupied(previous) => {
179                            if *previous.get() == op {
180                                trace!(
181                                    "         {wire:?} {op:?} change on {outlet_group:?} already done"
182                                );
183                            } else {
184                                debug!(
185                                    "         {wire:?} {op:?} change on {outlet_group:?} conflicting with {previous:?}. Blocked."
186                                    );
187                                return Ok(None);
188                            }
189                        }
190                    }
191                }
192            }
193        }
194    }
195    debug!("Translating {change:?} to patch");
196    let mut patch = TypedModelPatch::new(format!("{change:?}"));
197    let mut replaced_wires: HashMap<OutletId, OutletId> = HashMap::default();
198    let nodes_to_replace = changed_wires
199        .keys()
200        .flat_map(|outlets| outlets.iter().map(|o| o.node))
201        .chain(changed_ops.keys().copied())
202        .collect::<std::collections::HashSet<usize>>();
203    for node_id in model.eval_order()? {
204        let node = model.node(node_id);
205        if nodes_to_replace.contains(&node_id) {
206            let mut inputs = tvec!();
207            for (slot, orig) in node.inputs.iter().enumerate() {
208                let tgt = if let Some((outlet, alteration)) =
209                    rewired_scalar_input.get(&InletId::new(node_id, slot))
210                {
211                    let const_node = model.node(outlet.node);
212                    let mut value =
213                        const_node.op_as::<Const>().unwrap().val().clone().into_tensor();
214                    alteration.change_tensor(&mut value, false)?;
215                    let name = model.unique_name(&const_node.name);
216                    patch.add_const(name, value)?
217                } else {
218                    *replaced_wires
219                        .entry(*orig)
220                        .or_insert_with(|| patch.tap_model(model, *orig).unwrap())
221                };
222                inputs.push(tgt);
223            }
224            let op: Box<dyn TypedOp> =
225                changed_ops.get(&node_id).cloned().unwrap_or_else(|| node.op.clone());
226            let new_wires = patch
227                .wire_node(&node.name, op.clone(), &inputs)
228                .with_context(|| format!("wriring changed_op {op:?}"))?;
229            if new_wires.len() == 1
230                && patch.node(new_wires[0].node).op_is::<crate::ops::source::TypedSource>()
231            {
232                patch.inputs.insert(new_wires[0].node, node_id);
233            }
234            for (ix, w) in new_wires.iter().enumerate() {
235                replaced_wires.insert((node_id, ix).into(), *w);
236            }
237        } else {
238            for orig in &node.inputs {
239                if let Some(replacement) = replaced_wires.get(orig) {
240                    patch.shunt_outside(model, *orig, *replacement)?;
241                }
242            }
243        }
244    }
245    for output in model.output_outlets()? {
246        if let Some(replacement) = replaced_wires.get(output) {
247            unsafe {
248                patch.shunt_outside_unchecked(*output, *replacement)?;
249            }
250        }
251    }
252    let mut interface_change = tvec!();
253    for (ix, input) in model.input_outlets()?.iter().enumerate() {
254        if let Some(change) = changed_wires.get(&bound_outlets(*input)) {
255            interface_change.push((InOut::In(ix), change.clone()));
256        }
257    }
258    for (ix, output) in model.output_outlets()?.iter().enumerate() {
259        if let Some(change) = changed_wires.get(&bound_outlets(*output)) {
260            interface_change.push((InOut::Out(ix), change.clone()));
261        }
262    }
263    debug!("Patch ready for {change:?}");
264    Ok(Some((patch, interface_change)))
265}
tract_core/optim/change_axes.rs

tract_core/optim/
change_axes.rs
