pub fn insert_copies(graph: &SharedGraph, roots: &[NodeId]) -> Vec<NodeId>Expand description
Insert Op::Copy nodes so every op’s inputs live on the same
device as the op itself. Returns the rewritten roots.
This is the Phase-3.5 pass that lifts the Router from “explicit copies only” to “auto-insert where needed.”
§Device inference rules
For each node N, walked in topological order:
-
Target device = first match wins:
- If N has a placement hint via
graph.set_placement(n, d), used. - Else, inherit from N’s first input’s inferred device.
- Else (N is a Const or the first input has no inferred device
either), the target is
None— N is “placeless” and will be placed by its consumer’s demands.
- If N has a placement hint via
-
Input reconciliation: for each input I of N, if I’s inferred device differs from N’s target device, insert an
Op::Copy { target }node and redirect N’s input edge to the new Copy.
§What it doesn’t do
- Backward ops: the pass currently treats all ops the same way — no special reasoning about what device a backward node should run on. For pure inference this is fine; for training through the Router we’d want a reverse-mode pass that mirrors forward placement. Tracked as a follow-up TODO.
- Const placement lowering: if a Const flows into a single
consumer on device D, we emit
Copy(Const, D)instead of tagging the Const itself as on-device-D. A future cost-lowering pass can fuse the two — tracked as a TODO alongside Phase 4’s scheduler. - Redundant copy elision:
Copy(X, A) -> Copy(_, A)when X is already on A is dropped via the idempotent check in step 2, but there’s no pass that mergesCopy(X, A) -> Copy(_, B) -> Copy(_, A)(cross-device round-trips that cancel). Also future.