Skip to content

Fix AutoEnzyme OOB writes from short gradient buffers - #1405

Draft
ChrisRackauckas-Claude wants to merge 3 commits into
SciML:masterfrom
ChrisRackauckas-Claude:opt-enzyme-grad-size-check
Draft

ChrisRackauckas-Claude wants to merge 3 commits into
SciML:masterfrom
ChrisRackauckas-Claude:opt-enzyme-grad-size-check

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Oct 10, 2026 •

Copy link
Copy Markdown
Member

Summary

OptimizationEnzymeExt wrapped user gradient buffers as Enzyme.Duplicated(θ, res) with no size check. When res was shorter than θ, Enzyme wrote past the end of the buffer (auditor sentinel probe: one slot past G overwritten at n=4 and n=40 on 1.12 and 1.10; one run segfaulted in GC). This PR adds an O(1) size check that throws DimensionMismatch before Enzyme in every place a user-supplied buffer is Duplicated against θ, and also when user θ is paired with an internal cache sized from x (out-of-place grad/fg/hess/hv/cons paths and in-place BatchDuplicated hess/fgh/cons_j/cons_h/lag_h). The throw path is @noinline so the happy-path compare does not pull string interpolation into grad!.

Verification

Failing before / passing after (Julia 1.12.7; src stashed for before):

Before (has_check=false; short G segfaults instead of throwing):

VERSION=1.12.7
has_check=false
good_grad_ok G=[-90.99999999999999, 89.99999999999999]

[3513572] signal 11 (128): Segmentation fault
in expression starting at .../reg_grad_size.jl:22

After (throws DimensionMismatch for short G and wrong-length OOP θ):

grad PASS
fg PASS
oop PASS
VERSION=1.12.7 ALL PASS

Also confirmed on Julia 1.10.12. Sentinel probe after the fix: 0 slots overwritten.

Happy-path cost (BenchmarkTools minimum in ns, master vs PR):

master PR
1.10 n=4 grad 16.4 16.4
1.10 n=4 fg 24.9 24.9
1.10 n=40 grad 86.4 86.4
1.10 n=40 fg 126.5 126.5
1.12 n=4 grad 15.2–15.6 16.0
1.12 n=4 fg 23.7–24.1 24.0–24.5
1.12 n=40 grad 75.3–75.8 75.6–76.0
1.12 n=40 fg 124.7–126.6 124.6–126.6

Test-group tails

OPTIMIZATION_TEST_GROUP=AD on Julia 1.12:

Test Summary: | Pass  Total      Time
AD            | 1021   1021  14m47.8s
     Testing OptimizationBase tests passed

OPTIMIZATION_TEST_GROUP=QA on Julia 1.12:

Test Summary: | Pass  Broken  Total     Time
QA            |   20       1     21  2m19.4s
     Testing OptimizationBase tests passed

Out of scope / DI note

OptimizationEnzymeExt loads only when both Enzyme and ChainRulesCore are loaded (weakdeps pair). With only using OptimizationBase, Enzyme, AutoEnzyme goes through DifferentiationInterface, and a short G there can still write one slot past the end. That DI path is out of scope for this PR.

What I did NOT verify

  • Full GROUP=All / root monorepo suite
  • Opening a separate DI-side issue for the short-G OOB when CRC is not loaded

What a reviewer should push back on

  • Whether every BatchDuplicated/x-cache site listed above needs the check vs. documenting that θ must match u0
  • Whether OptimizationBase needs a patch version bump for this bug fix alone

Please ignore until reviewed by @ChrisRackauckas.

Risk assessment

Independent review: pending

🤖 Generated with Cursor Agent 2026.10.01-14929f9 (model: unknown (Cursor auto)), transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-enzyme/log.txt on amdci2; orchestrated by Claude Code (claude-opus-5-5[1m]) https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

Independent review: Devin Fusion (fusion-claude-opus-5-5-high-sidekick-swe-2-medium) rated it low, verdict MERGE: #1405 (comment)

ChrisRackauckas and others added 3 commits October 10, 2026 00:58
Enzyme.Duplicated(θ, res) writes with θ's shape; a too-short user buffer
was memory-unsafe (sentinel clobber / GC segfault). Throw DimensionMismatch
before autodiff in grad!/fg! and other user-buffer Duplicated paths.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Cursor Agent <noreply@cursor.com>
Agent-Harness: Cursor Agent 2026.10.01-14929f9
Agent-Model: unknown (Cursor auto)
Agent-Session: local session, transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-enzyme/log.txt on amdci2
Move DimensionMismatch construction to @noinline _throw_dup_size so the
O(1) size compare does not inflate correct-size grad! cost, and reject
wrong-length θ against x-sized internal Enzyme caches (OOP grad/fg/hess
and BatchDuplicated paths).

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Cursor Agent <noreply@cursor.com>
Agent-Harness: Cursor Agent 2026.10.01-14929f9
Agent-Model: unknown (Cursor auto)
Agent-Session: local session, transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-enzyme/log.txt on amdci2
Use a neutral expected/got DimensionMismatch message so _check_duplicated_size(x, θ)
does not mislabel the user's θ as the shadow buffer.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Cursor Agent <noreply@cursor.com>
Agent-Harness: Cursor Agent 2026.10.01-14929f9
Agent-Model: unknown (Cursor auto)
Agent-Session: local session, transcript /home/crackauc/sandbox/goals/performance/jobs/opt/opt-enzyme/log.txt on amdci2
@ChrisRackauckas

Copy link
Copy Markdown
Member

🤖 Automated comment from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris.

Independent review (Devin CLI 3000.11.3, model fusion-claude-opus-5-5-high-sidekick-swe-2-medium): MERGE, risk low.

Full review

VERDICT: MERGE
RISK: low

Blocking findings

None.

The fix is correct. Every check compares sizes before Enzyme pairs two buffers whose memory has to line up. I could not find a site where the check rejects a call that used to be valid. Enzyme.Duplicated/BatchDuplicated already require both arguments to have the same type, so comparing size rather than length only rejects same-type arrays with different shapes, and those were never supported. Shadow caches are built from x as Vectors, so a θ of a different type already raised a MethodError on master. I confirmed the bug and the fix by running them (see "What I ran").

Non-blocking findings

  1. Most of the new checks have no regression test (confirmed by reading lib/OptimizationBase/test/AD/adtests.jl:1539-1553). The new testset covers in-place grad and fg with a short G, and out-of-place grad with a long θ: 3 of about 33 _check_duplicated_size calls. If the checks in hess, fgh!, hv!, cons_j!, cons_vjp!, cons_jvp!, cons_h! and lag_h! (both IIP and OOP) were deleted, no repo test would fail. My probe ran 26 mismatch cases across those paths and every one threw DimensionMismatch. It would still be worth folding a looped version of those cases into the testset.
  2. The error message does not say which argument is wrong (ext/OptimizationEnzymeExt.jl:171-177). It reads "expected size (2,) for Enzyme Duplicated, got (1,)" without naming the gradient buffer, θ, or v. This is a cosmetic issue.
  3. Pre-existing bugs found while probing. The PR does not touch these lines, and each one behaves the same on master. They should be filed separately and are not reasons to hold this PR:
    • fgh! returns a zero gradient (confirmed by running on master and on the PR, Julia 1.10 and 1.12). IIP fgh!(G, H, θ) leaves G == [0, 0] for Rosenbrock at θ = [0.5, 0.7], where the expected value is [-91, 90]. OOP fgh!(θ) also returns g == [0, 0], though H is correct. The cause is that G/G_fgh is passed as the primal of BatchDuplicatedNoNeed (ext/OptimizationEnzymeExt.jl:341, :757), so the first-order result is never written.
    • IIP cons_vjp!(res, θ, v) zeroes the caller's v (confirmed by running). v is the reverse-mode output shadow in Duplicated(cons_res, v) (:477), and Enzyme consumes it.
    • OOP hess/lag_h! rebind the captured variable H, which breaks hv! (confirmed by running on the PR branch; the code is unchanged from master). The H = Matrix(...) at :702 and :936 assigns to the outer H = zero(x) at :777. After hess(θ) has been called, hv!(θ, v) throws MethodError: no method matching Duplicated(::Vector{Float64}, ::Matrix{Float64}).
  4. Two CI jobs are red, both unrelated to this PR. downgrade-sublibraries / test (lib/OptimizationReactant) fails in the same way on other branches (runs 38059942878 and 38041080820). The ModelingToolkit downstream job ended after 30 s, during checkout or cleanup.
  5. There is no OptimizationBase version bump. The PR body raises this itself, and it is the maintainer's decision.

What I ran

Logs are in scratch/. Master code came from git archive origin/master lib/OptimizationBase, using merge-base aac9e45.

  • scratch/probe.jl, a probe I wrote, run against the PR and against master on Julia 1.12.7 and 1.10.12 (probe2_{pr,master}_{1.12,1.10}.log). It does four things:
    • Happy path vs analytic answers. It checks Rosenbrock gradient, Hessian and Hv, plus constraint J, Jᵀv, Jv, the constraint Hessians and the Lagrangian Hessian (full and lower-triangle forms), for IIP and OOP, against closed-form derivatives. It also checks grad, fg and hess for a 2×2 matrix-shaped u0. Results were 12/12 IIP, 9/9 OOP and 3/3 matrix on both branches, so the checks do not reject valid calls.
    • Sentinel test for the out-of-bounds write. A length-(n−1) G is unsafe_wrapped over a buffer filled with 12345.0. A first version used NaN, which hid the write because Enzyme accumulates with +=. Results:
      • master, n=4, both Julia versions: grad and fg threw nothing and wrote 1 slot past the end of the buffer.
      • PR, n=4 and n=40, both Julia versions: DimensionMismatch in every case and 0 slots written past the end.
    • Mismatch cases on the PR. 16/16 IIP and 10/10 OOP threw DimensionMismatch, covering every check site except OOP cons_vjp/cons_jvp. Those two paths are gated on f.cons_vjp == true and are effectively unreachable.
  • The PR's new testset, extracted verbatim.
    • PR on 1.12: passed 3/3.
    • master on 1.12: optf.grad(Gshort, ...) reported "No exception thrown", and the process then died with SIGSEGV (signal 11). That is a fail before the fix and a pass after it.
  • test/AD/enzyme_lagrangian_hessian.jl with the PR code on 1.12: 134/134 passed.
  • test/AD/adtests.jl with the PR code on 1.12: 840/840 passed, including the new testset.
  • gh pr checks: OptimizationBase AD (1 and lts), QA, Core, Runic and typos all pass. I checked the two red jobs against other branches (finding 4).

What I did not verify

  • The happy-path benchmark table in the PR body. I did not rerun BenchmarkTools. From reading the code, the check is a single size tuple compare with an @noinline throw.
  • The PR's claim that without ChainRulesCore loaded, a short G still writes out of bounds through the DifferentiationInterface path.
  • The full OPTIMIZATION_TEST_GROUP=AD group (I ran only the two files above, and skipped dual_tolerant_tests.jl and cvxtest.jl), and the QA group locally. CI shows both green.
  • Nested-Enzyme (outer AD through cons_j!) with the extra x captured in the let block at :432-433. The only evidence I have is that enzyme_lagrangian_hessian.jl passes.
  • There is no linked issue or tracking issue (closingIssuesReferences is empty), so I could not check fit with sibling work.

🤖 Posted by an AI agent — harness: Devin CLI 3000.11.3 (review), Claude Code 2.1.285 (posting) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium
Conversation: local Claude session 3c311569-dda6-4687-bc9f-65febb212c33 (fleet master)

This branch has not been deployed

No deployments
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.

2 participants