Skip to main content

insert_layout_fixups

Function insert_layout_fixups 

Source
pub fn insert_layout_fixups<F>(
    graph: &mut Graph,
    roots: &[NodeId],
    kernel_accepts_strided: F,
) -> usize
where F: Fn(NodeId) -> bool,
Expand description

Phase 2.2 of the picker arc — insert Op::Contiguize before any kernel-bearing node whose chosen kernel doesn’t accept strided inputs but whose live input layout is non-contiguous (strided, offset, or both).

§When to call

After Phase 2.1’s insert_cross_device_copies and after the picker has committed to per-node kernel choices. The picker’s commitment determines the KernelCaps::strided_input flag for each node; this pass reads that flag via the kernel_accepts_strided callback and inserts layout-fixups where the flag is false and the input is strided.

§Callback contract

kernel_accepts_strided(consumer_id) -> bool returns true if the kernel resolved for that node tolerates strided / offset inputs (i.e. its KernelCaps::strided_input == true). The callback is opaque to this pass — typical wiring reads the picker’s commit-to-winner alternative set and queries the winner’s caps. When false, this pass inserts Op::Contiguize on every non-contiguous input.

§Skips

  • View ops (Transpose, Permute, BroadcastTo, Slice, Unsqueeze, Squeeze, Flip) — they preserve strided layouts intentionally.
  • Structural ops (Const, Release, Reshape, Contiguize, Copy, Move, Alloc, ZeroFill, WriteSlice) — their executor arms handle strided naturally or are not kernel- bearing.

§CSE

When two consumers share a non-contiguous input AND both kernels reject strided, one Op::Contiguize node serves both. The per-pass cache is keyed on producer NodeId (no target-device dimension — Contiguize is device-inheriting).

§Idempotence

Running this pass on an already-rewritten graph adds zero new nodes: an inserted Op::Contiguize has a contiguous + zero-offset output layout, so the subsequent fixup check (is_contiguous && start_offset == 0) succeeds and the input is skipped.

Returns the number of Op::Contiguize nodes inserted.