Unseed only the chunk, and split unseed! out of seed! - #821
Unseed only the chunk, and split unseed! out of seed!#821ChrisRackauckas-Claude wants to merge 3 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #821 +/- ##
==========================================
- Coverage 90.74% 90.66% -0.09%
==========================================
Files 11 11
Lines 1070 1071 +1
==========================================
Hits 971 971
- Misses 99 100 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
CI: all 19 Julia 1 / lts / min-patch jobs green. The 6 "Julia pre" failures are the pre-existing Julia 1.13-rc1 + JET issue that also fails on master — root-cause analysis in #816 (comment). This commit introduces no new failures relative to the #816 branch it stacks on. |
|
Marked as draft since it depends on #816. Seemingly could be decoupled if supposed to be reviewed more urgently. |
The chunk-mode unseed call seed!(xdual, x, i) only needs to clear the N-wide chunk seeded at i: the rest of the array is zeroed up front and every other chunk clears itself. ForwardDiff 0.10 wrote exactly N elements here; the 1.x rewrite made it write from i to the end of the array, i.e. O(n^2/2N) redundant dual writes per chunked gradient/jacobian sweep (~40 GB of memory traffic for gradient! of 100000 elements at chunk 12). Write at most N elements starting at index. The 3-arg seed! form has no other callers in the package. gradient! of sum(abs2, x) at chunk 12: n=1000 502 -> 279 us, n=100000 5.39 -> 2.99 s. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019GcNzbNzaHCqTm4W14eKhN
bacc1a2 to
5e17c90
Compare
|
Unstacked: force-pushed Re-verified locally on the unstacked branch: full |
| # Writes at most N elements starting at `index`: chunk mode only ever needs to | ||
| # clear the N-wide chunk it just seeded, so writing through to the end of the | ||
| # array would be O(n) redundant work per chunk (O(n^2) per sweep). |
There was a problem hiding this comment.
This comment is confusing - albeit seemingly correct - since the function is not actually used to "seed" anything but to unseed existing seeds (to replace by zero partials). Can we decouple the unseeding and seeding calls? E.g. use unseed!(duals, x, index) instead of seed! with an implicit zero(Partials(...)) argument?
There was a problem hiding this comment.
Agreed — done in b8a9b17.
seed!(duals, x) and seed!(duals, x, index) are now unseed!(duals, x) and unseed!(duals, x, index), and I dropped the seed::Partials argument from them rather than keeping it as an implicit zero: nothing in the package ever passed it (the only caller of the explicit-seed form was AllocationsTest, which now exercises unseed! instead). The NTuple{N,Partials} methods keep the seed! name and are otherwise untouched, so the two operations are now separate functions rather than one function distinguished by a default argument.
I renamed the whole-array form too, not just the indexed one flagged here — seed!(xdual, x) before the chunk loop and seed!(ydual, y) for the output buffer are unseeding as well, so leaving that one as seed! would have kept exactly the confusion you're pointing at. Happy to narrow it back to the 3-arg method if you'd rather keep the diff minimal.
Verified locally after the rename: full Pkg.test() passes (9134/9134, Julia 1.12), including the allocation tests, and the benchmark numbers are unchanged (unseed sweep at n=100000: 2.49 s on master, 0.29 s here).
Review feedback: the zero-Partials default of seed! made the unseeding calls read as seeding. Give them their own name and drop the seed argument, which was never passed for those two methods (the only caller of the explicit-seed form was the allocation test). seed!(duals, x) -> unseed!(duals, x) seed!(duals, x, index) -> unseed!(duals, x, index) The NTuple-of-seeds methods are unchanged. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019GcNzbNzaHCqTm4W14eKhN
devmotion
left a comment
There was a problem hiding this comment.
LGTM. Can you update the version number?
Summary
Standalone against
master(previously stacked on #816; now independent of it — the twotouch the same function and will need a trivial conflict resolution whichever lands second).
In ForwardDiff 0.10, the chunk-mode "unseed" call
seed!(duals, x, index, seed)wroteexactly the
Nchunk elements starting atindex. The 1.x rewrite changed it to writefrom
indexto the end of the array:Chunk mode calls this after every chunk (
seed!(xdual, x, i)inchunk_mode_gradient_expr/chunk_mode_jacobian_expr) purely to clear the chunk itjust seeded — the rest of the array is already unseeded (everything is zeroed once
up front, and every other chunk clears itself). So writing to the end does
Σ(n − i) ≈ n²/2N redundant dual writes per gradient/jacobian sweep. At n = 100000 with
chunk 12 and
Dual{…,Float64,12}(104 bytes each) that is ~40 GB of memory traffic.This restores the 0.10 behavior: write at most
Nelements starting atindex(
Iterators.take(…, N)). The 3-argseed!form has no other callers in the package.Second commit, from review: the two
seed!methods whoseseed::Partialsargumentdefaulted to
zero(Partials)were only ever called to unseed, so they are nowunseed!(duals, x)andunseed!(duals, x, index)with that argument dropped (it wasnever passed outside the allocation test). The
NTuple{N,Partials}seeding methods keepthe
seed!name and are unchanged.Benchmarks
Julia 1.12,
gradient!off(x) = sum(abs2, x)with default chunk 12, median of 11 runs:gradient!n=1000gradient!n=10000gradient!n=100000The remaining unseed cost is the
Iterators.drop(structural_eachindex(…), index-1)walk,which is still O(index) per chunk; #816's dense fast path removes that separately.
Tests
No new tests: chunk-vs-vector-mode consistency (including non-divisible chunk sizes and
UpperTriangular/Diagonalinputs) is covered by the existing suite, which passeslocally (
9134/9134on Julia 1.12). Verified additionally that chunked results matchfull-chunk references for n ∈ {5, 12, 13, 24, 25, 100} × chunk ∈ {1, 3, 5, 12}, for
gradient,jacobian, and the mutatingjacobian(f!, y, x)form, plusUpperTriangular/
Diagonalstructural inputs and aBigFloat(non-isbits) input.Note
Opened as a draft by an agent on behalf of @ChrisRackauckas. Please ignore
until reviewed by @ChrisRackauckas.
🤖 Generated with Claude Code