Skip to content

Orbit: Bypass handling for residual connections - #3383

Open
edwjchen wants to merge 2 commits into
google:mainfrom
edwjchen:orbit-bypass
Open

edwjchen wants to merge 2 commits into
google:mainfrom
edwjchen:orbit-bypass

Conversation

@edwjchen

Copy link
Copy Markdown
Collaborator

Stacks on #3194.

A residual connection is a fork that feeds a deep "main" path and a shallow path which rejoin at an addition. No single value is live across that span, so the region has no single-input single-output cut and SISO partitioning would leave it as one monolithic ILP. When the two rejoined operands differ in multiplicative depth by at least bypass-depth-threshold (Orbit's bpsdepth, default 15), the region is instead split and solved separately.

How it works

  • Detection (detectBypass): finds the fork/join in a region, e.g., an output addition with two secret operands from distinct producers whose mul-depth gap clears the threshold.
  • Split solve (solveBypass): the deep main path is solved for every enumerated output level to build a QBP table; the shallow path plus the join is then solved once per realized main-output state under the region's own output boundary; the cheapest main + bypass combination is kept.
  • Runs on every region (solveRegion dispatcher)
  • Residual-join carve-out: the join addition is kept out of any squashed addition tree (via per-op multiplicative depth in OpGrouping) so its two-predecessor structure stays visible to detection.
  • Contested fork: the fork is pinned to the region's DP-enumerated input state and each path rescales it independently per edge, so a rescale common to both paths is charged once (absorbed into the previous partition's output) rather than duplicated.

To make bypass regions first-class, the partition driver was generalized from a contiguous [begin, end) range to an explicit op-set (groupIds), so main and bypass can be arbitrary subsets of a region.

Default

On by default at threshold 15, matching Orbit; set bypass-depth-threshold=0 to disable. Existing circuits without a large-gap residual are unaffected.

Port Orbit's three ILP-reduction techniques: addition-tree squashing and
structural compression (auto_compress) merge equivalent ops into groups that
share one set of ILP variables (exact, since op costs are linear in level),
and single-input single-output partitioning solves large circuits piecewise
under enumerated boundary states with a DP stitch.
Port Orbit's bypass (handle_bypass): a fork feeding a deep main path and a
shallow path that rejoin at an addition has no single-input single-output
cut, so the region is split into main and bypass and solved separately
instead of as one monolithic ILP. The main path is solved for every
enumerated output level to build a transfer table, the shallow path plus the
join is solved once per main-output state, and the cheapest combination is
kept. Detection runs on every region — the whole body and each partition of
the DP — so mid-circuit residuals are handled; the residual-join addition is
kept out of squashed addition trees so its structure stays visible. On by
default (bypass-depth-threshold=15, matching Orbit); set to 0 to disable.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant