Skip to main content

insert_copies

Function insert_copies 

Source
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:

  1. Target device = first match wins:

    • If N has a placement hint via graph.set_placement(n, d), use d.
    • 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.
  2. 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 merges Copy(X, A) -> Copy(_, B) -> Copy(_, A) (cross-device round-trips that cancel). Also future.