Conversation
ActivationRematerializationPass emitted a quadratic number of operations for a deep producer chain. Materializing a tagged op clones it once per surviving consumer, and a tagged CONSUMER contributes one consumer per clone it will itself produce, so the counts compound along a chain. The greedy selection has no term for this — it prices each candidate by its own recompute cost alone — and happily chose such a plan. Measured on a 2000-deep chain with every intermediate live to a sink: 4,001 ops became 2,001,002 in 16.4s. Three parts: 1. projectCloneCount computes a plan's clone count EXACTLY before any IR is built, mirroring the reverse materialization walk. Verified exact at depths 4/8/16/24/32/40/48: emitted ops always equal input + projected - selected originals erased. This is what makes it usable as a gate rather than a post-hoc count. 2. The plan is trimmed to fit --max-clone-expansion x |ops| (default 8), warning REMAT_PLAN_CLONE_BOUND with what it gave up. Trimming drops the op nearest the MIDDLE of the tagged chain, found via cloneCounts (which rises monotonically along a chain), so one chain of length K becomes two of K/2 and the projection falls from ~K^2/2 to ~K^2/4 per drop — the segmentation real activation checkpointing performs. 3. A dropped op's interval goes back into the difference array, so the reported peak-after describes the plan actually emitted. On the 2000-deep chain: 33,842 ops in 2.4s with peak cut 93.5%, against 2,001,002 ops in 16.4s for a 99.8% cut. This is a TRADE, not a strict win — ~6 points of peak reduction for a 59x smaller function — and --max-clone-expansion=0 restores the unbounded behavior. Choosing the midpoint matters: dropping the maximum-clone op instead only peels the chain's downstream end, giving 35,627 ops for a mere 12.5% peak cut. Explicit tessera.recompute markers stay authoritative and are never trimmed (verified: same chain with explicit markers and a tight bound is untouched, 91 ops, no warning) — the bound applies only where this pass chose the plan itself. lit 439/439. Remat-related unit tests 268 passed / 5 skipped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1db4036d35
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| func->emitWarning() | ||
| << "REMAT_PLAN_CLONE_BOUND: the budget-selected recompute plan " |
There was a problem hiding this comment.
This shared compiler-pass change introduces a new diagnostic and IR planning attribute, but the commit leaves docs/audit/backend/{apple,nvidia,rocm,x86}/todo.md untouched. Assess the impact for each backend and record whether parity is validated, follow-up is required, or the change is not applicable so the architecture queues do not drift.
AGENTS.md reference: AGENTS.md:L81-L85
Useful? React with 👍 / 👎.
| droppedOps.insert(*split); | ||
| selected.erase(split); |
There was a problem hiding this comment.
Whenever the clone bound trims at least one candidate, selected.erase(split) changes the emitted plan but selectedCost still includes every dropped candidate's recompute cost from the greedy loop. As a result, tessera.remat_selected_cost_ns describes the rejected pre-trim plan while the peak, clone-count, and selected-count attributes describe the emitted plan, misleading any planner or telemetry consumer comparing those values; recompute or decrement the cost as candidates are dropped.
Useful? React with 👍 / 👎.
1 participant
Stacked on #640 (it builds on that PR's peak-scan fix in the same pass).
The defect
ActivationRematerializationPassemits a quadratic number of operations for a deep producer chain. Materializing a tagged op clones it once per surviving consumer, and a tagged consumer contributes one consumer per clone it will itself produce — so the counts compound along a chain. The greedy selection has no term for this:estimateRecomputeCostprices each candidate by its own cost alone, so recomputing element i of a chain looks cheap when it actually costs the whole prefix.Measured on a 2000-deep chain with every intermediate live to a sink: 4,001 ops became 2,001,002 in 16.4s.
Confirmed the mechanism directly before fixing it — output op count at depths 4/8/16/32 was 15/45/153/561, quadrupling per doubling.
The fix
projectCloneCountcomputes a plan's clone count exactly, before any IR is built, by mirroring the reverse materialization walk. Verified exact at depths 4/8/16/24/32/40/48: emitted ops always equalinput + projected − selected originals erased. That exactness is what makes it a usable gate rather than a post-hoc count.--max-clone-expansion × |ops|(default 8), warningREMAT_PLAN_CLONE_BOUNDwith what was given up. Trimming drops the op nearest the middle of the tagged chain — located viacloneCounts, which rises monotonically along a chain — so one chain of length K becomes two of K/2 and the projection falls from ~K²/2 to ~K²/4 per drop. That is the segmentation real activation checkpointing performs.peak_afterdescribes the plan actually emitted.Results on the 2000-deep chain
This is a trade, not a strict win — about 6 points of peak reduction for a 59x smaller function — and
--max-clone-expansion=0restores the unbounded behavior for a caller who wants the last of the memory.The third row is why the midpoint rule matters and is worth stating: my first attempt dropped the maximum-clone op, which only peels the chain's downstream end and leaves a shorter contiguous chain. It cut code size fine but achieved almost no memory benefit. I caught it by checking
projected / selected ≈ selected / 2, the signature of a single chain rather than segments — the comment claiming "segmentation" was wrong until the rule was changed to match it.Scope
Explicit
tessera.recomputemarkers stay authoritative and are never trimmed — verified: the same chain with explicit markers and a tight bound is untouched (91 ops, no warning). The bound applies only where this pass chose the plan itself.lit439/439 (new fixtureremat_clone_expansion_bound.mlirpins both the unbounded projection and the trimmed plan). Remat-related unit tests 268 passed / 5 skipped.REMAT_PLAN_CLONE_BOUNDregistered indiagnostic_codes.pyandpass_metadata.py.🤖 Generated with Claude Code