Skip to main content

Module inline_for_partition

Module inline_for_partition 

Source
Expand description

Selective function inlining for the partition stage.

Modules don’t drive network boundaries any more — wire ops do. Three classes of functions get inlined at every CALL site before [crate::partition_by_wire_ops] runs:

  1. Wire-touching functions — any function whose transitive closure contains an ai.bytesandbrains.wire op. Wire ops must live at the top level so the partitioner’s reachability walk cuts the graph cleanly; a wire hidden behind a CALL fragments the partition boundary.

  2. Pure-ONNX functions — any function whose transitive closure is entirely ai.onnx.*. Inlining surfaces each NodeProto at the top level so the engine can route per-op against the bound Backend’s Contract methods (add, matmul, …) without an intervening CALL indirection.

  3. Single-call functions — any function CALLed from exactly one site across the whole call graph. Keeping it as a FunctionProto saves no memory (the body would appear once either way) but adds an indirection at dispatch time, so we inline it eagerly.

Functions called from multiple sites (and not in classes 1 or 2) survive as FunctionProto — the body appears once, callers reference it via CALL nodes.

The root function (model.functions[0]) is always preserved regardless of its body’s classification — it’s the entry point the compiler partitions on.

Functions§

inline_for_partition
Inline every wire-touching or pure-ONNX function at every CALL site. Iterates to a fixed point because inlining a wire-touching callee may itself reveal a new wire-touching caller. Returns the total number of CALL replacements performed.