Seed dense arrays of isbits duals without scalar indexing (fixes GPU jacobians) - #816
ChrisRackauckas-Claude wants to merge 1 commit into
Conversation
|
Because the broadcasts were purposefully removed due to performance things, we should at least recover the GPU functionality by handling AbstractGPUArray. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #816 +/- ##
==========================================
+ Coverage 90.68% 90.94% +0.25%
==========================================
Files 11 11
Lines 1052 1082 +30
==========================================
+ Hits 954 984 +30
Misses 98 98 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
It's difficult to verify this claim given that the As an alternative to an extension, I wonder whether special-casing |
18a692a to
35d4f92
Compare
|
Good suggestion — I benchmarked it, and the On the interface concern: dispatching on duals isa DenseArray && isbitstype(V) && !Base.has_offset_axes(duals, x)so nothing needs to be assumed about GPU element types either — non-isbits
|
| method | writes | loop (master) | broadcast | map! |
|---|---|---|---|---|
seed!(duals, x, seed), n=1000 |
all n | 4.1 μs | 2.7 μs | 4.1 μs |
seed!(duals, x, seeds), n=1000 |
first N | 56 ns | 56 ns | 52 ns |
seed!(duals, x, 500, seed), n=1000 |
index:end | 3.0 μs | 1.4 μs | 2.0 μs |
seed!(duals, x, 500, seeds), n=1000 |
N at index | 458 ns | 1.6 μs (7 allocs) | 62 ns |
seed!(duals, x, 50000, seeds), n=100000 |
N at index | 40 μs | 1.6 μs (7 allocs) | 68 ns |
So the answer to "map! or broadcast?" is: both, per method (bold = what the PR now uses).
- For the full-array write (method 1), broadcast wins; on 1.10
map!is actually ~25% slower than the loop there, so the PR uses broadcast. - For method 3 (index:end), broadcast is fastest but its dotview allocates 112 bytes under
--check-bounds=yeson Julia 1.10 (whichPkg.testuses, andtest/AllocationsTest.jlrightly rejects), so the PR usesmap!over views — allocation-free on both versions and still 1.5× faster than the loop. - For the chunk writes (methods 2 and 4),
map!wins — but the seedsNTuplecan't be amap!source (there's nomap!(f, dest, ::Tuple)method), and slicing it at runtime (seeds[1:chunksize], what both 0.10 and the extension in the first iteration of this PR did) is what costs the 1.6 μs + allocations in the broadcast column.map!over the index range with a closure over the tuple avoids both, and GPUArrays'map!handles the range argument fine. - The last row is the real story for chunk mode: the structural path pays
Iterators.drop(structural_eachindex(duals, x), offset)— O(index) per chunk, so O(n²/N) per full chunked jacobian/gradient sweep. The dense path indexes the chunk directly.
On the 2018 revert (#354) that removed broadcast for performance: the regression jrevels measured there was pre-Julia-1.0 broadcast plus an allocating x[dual_inds] slice; with views and Ref none of these forms allocate today.
End-to-end (master → this PR, Julia 1.12.4)
gradient!n=1000 (chunk 12): 464 μs → 344 μsgradient!n=100000 (chunk 12, median of 7 runs): 4.80 s → 3.96 sjacobian!n=100: 63 μs → 46 μsjacobian!n=1000: 5.92 ms → 4.84 msgradient!n=10 (vector mode): 141 ns → 158 ns — the one regression;seed!itself measures identical (46 ns both), so it's an inlining side effect of the added branch.
One more thing that turned up while measuring: the 1.x rewrite also changed the chunk-mode "unseed" call seed!(xdual, x, i) to write from i to the end of the array each chunk — 0.10 wrote exactly the N chunk elements — which is O(n²) dual writes per sweep (~40 GB of memory traffic for the n=100000 gradient above). Restoring the narrow unseed on top of this PR takes the n=100000 gradient! from ~3.9 s to ~2.4 s. I've kept that out of this PR to keep it focused; follow-up PR coming.
JLArray tests (vector/chunk mode, jacobian!, f!, matrix and view inputs) now exercise the main-package code path directly; the extension and weak dep are gone.
|
CI status: all 19 Julia 1 / lts / min-patch jobs (3 OSes × NaN-safe on/off) plus Documentation are green. The 6 "Julia pre" jobs fail, which needs a note:
The chunk-unseed O(n²) follow-up mentioned above is now open as #821 (stacked on this branch). |
|
@maleadt maybe should comment here on what is expected to be public and used. |
35d4f92 to
7479a2a
Compare
|
Rebased onto current |
7479a2a to
c91c578
Compare
|
All existing human review comments are addressed on current head
The post-review full suites pass 9309/9309 tests on Julia 1.10.12 and 9327/9327 on Julia 1.12.6, including allocation and QA checks. CI for The downstream conclusion remains release-gated: do not revert SciML/DeepEquilibriumNetworks.jl#243 until DeepEquilibriumNetworks requires a ForwardDiff release containing this fix. After that, reverting is appropriate only if the package intends to restore ForwardDiff GPU Jacobians and an actual CUDA steady-state-adjoint regression passes; otherwise keeping |
ForwardDiff scalar-indexes GPU arrays while seeding Jacobian chunks, and the chunk-tail zeroing added in 1.4.4 introduced the same failure in seed_zero_partials!. Use broadcast and map! over contiguous views for dense arrays with isbits values while preserving the structural fallback for wrappers, offset axes, and non-isbits values. Add JLArray regressions for both seeding paths, including the 1.4.4 tail clear, vector and chunk modes, mutating Jacobians, matrices, and views. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Co-Authored-By: Codex <noreply@openai.com> Original-Agent-Harness: Claude Code (version unknown) Original-Agent-Model: Claude Fable 5 Original-Agent-Session: https://claude.ai/code/session_01Vx7zQ96NYk4VV4ML2s3kAC Agent-Harness: Codex CLI 0.150.1 Agent-Model: unknown Agent-Session: 01a04603-0ac8-7570-9713-851acf5b8f5d
c91c578 to
8697dff
Compare
|
CI attempt 1 for tree I reproduced the same four values and failures byte-for-byte on clean The clean-master investigation is now complete:
That unrelated test correction is intentionally not mixed into this GPU-seeding PR. Current PR head |
|
I checked whether the dense On current head
Chunk zeroing was effectively identical. For a 100,000-element full-buffer zero, broadcast, The dense However, the structural/indexing fallback cannot be removed:
Conclusion: keep the current split. Dense, one-based arrays with isbits values should use the Commands used: ~/.juliaup/bin/julia +1.10 --project=../bench_env_110 ../bench_env/map_vs_index.jl
~/.juliaup/bin/julia +1.12 --project=../bench_env ../bench_env/map_vs_index.jlLinks |
|
@devmotion I think this is good for another review now. |
What changed and why
ForwardDiff GPU Jacobians currently fail because the seeding helpers scalar-index GPU arrays. An independent clean-checkout bisect found two regression points:
fce7f76cf8885ded3e150f11645c772961ce8dc5(released in v1.0.1) replaced broadcast seeding with scalarseed!loops.a337ee6658ed2a26bfa7251f1580da59c0b36625(merged in PR 821 and released in v1.4.4) added the chunk-tailseed_zero_partials!call, which introduced a second scalar-indexing path.The older
seed!failure normally masks the newer tail-clear failure. This PR therefore fixes both paths: dense arrays with isbits values use broadcast for full-buffer writes andmap!over contiguous views for chunk writes. A concrete isbits callable carries the seed tuple and offset intomap!, avoiding captured closures. Structural wrappers, offset axes, and non-isbits values retain the existing structural fallback.The runtime implementation uses only Base array interfaces and has no GPUArraysCore dependency or
AbstractGPUArraydispatch. JLArrays is added only as a test dependency; it is MIT-licensed, matching ForwardDiff's license.Downstream fallout
DeepEquilibriumNetworks PR 243 works around the v1.4.4 failure by selecting
autodiff=falsefor GPU steady-state adjoints. Do not revert that workaround before DeepEquilibriumNetworks requires a ForwardDiff release containing this PR: reverting earlier restores the CuArray failure.After such a ForwardDiff release is required, the regression-specific need for PR 243 is gone. Whether to revert it is then a policy choice:
autodiff=falsedefault.In short: treat PR 243 as release-gated cleanup, not an immediate revert.
Regression evidence
The final GPU-like test file was applied to clean
masterand run with scalar indexing disabled:The same test on this PR passes:
The focused command was:
Full verification
Both full runs include allocations and JET QA. These checks also exited successfully with no output:
~/.juliaup/bin/julia +release -m Runic --check test/GPUArraysTest.jl test/runtests.jl typos Project.toml src/apiutils.jl test/GPUArraysTest.jl test/runtests.jl git diff --checkFresh Julia 1.12.6 process-local medians against current
masterafter PR 821:gradient!, n=1000jacobian!, n=100jacobian!, n=1000Small n=10 gradient timings varied between process runs, so no directional claim is made for that case.
The review-requested replacement of captured closures with the concrete callable was also compared in fresh Julia processes. Against the preceding PR head, medians were 223.9 µs vs 222.5 µs for
gradient!at n=1000, 38.6 µs vs 27.2 µs forjacobian!at n=100, and 3.86 ms vs 3.30 ms forjacobian!at n=1000. Allocation tests in the full suite remain green.CI
All 27 checks on current head
8697dffpass, including the full Linux/macOS/Windows matrix across minimum-patch, LTS, Julia 1, and prerelease versions with NaN-safe mode enabled and disabled: https://github.com/JuliaDiff/ForwardDiff.jl/actions/runs/33137219763. Codecov reports 90.94% project coverage (+0.25% frommaster) and covers every modified line: https://app.codecov.io/gh/JuliaDiff/ForwardDiff.jl/pull/816.The preceding tree-identical CI attempt had one randomized
gamma_inctolerance failure on Windows Julia 1.12.7. The exact values reproduce on cleanmasterwith that job's seed and do not involve this PR's code; the failing job and clean-master reproduction are documented at #816 (comment). No assertion was loosened or silenced.Not verified locally
Note
This remains a draft opened by an agent on behalf of @ChrisRackauckas. Please ignore it until reviewed by @ChrisRackauckas.
Links
🤖 Generated with Codex CLI 0.150.1 (model: unknown; session: 01a04603-0ac8-7570-9713-851acf5b8f5d)
Risk assessment
seed!andseed_zero_partials!run on everygradient,jacobianandhessiancall. The new fast path covers anyDenseArrayinput with an isbits element type and no offset axes. That includes plainVector/Matrixinputs, so almost every CPU user of ForwardDiff hits the changed code, not just GPU users. Anything that breaks would show up in ForwardDiff's large downstream set (SciML, Optim, Turing and others). Non-isbits values, offset axes and structured wrappers likeDiagonalandUpperTriangularstill take the existing loop.8697dff: all 27 checks pass (https://github.com/JuliaDiff/ForwardDiff.jl/actions/runs/33137219763). That is 24 test jobs (Julia {min-patch, lts, 1, pre}×{ubuntu, macOS, windows}-latest× NaN-safe on/off), plusDocumentation,codecov/patchandcodecov/project. Codecov reports every modified line covered, with project coverage up 0.25%.a3c0f4f: the same matrix passes in the push run, including theJulia prejobs. No check fails on head, so there is nothing pre-existing to compare.gamma_incfailure on the previous tree-identical attempt was reproduced on clean master with that job's seed, and no assertion was changed.eigenandeigvals#757,eigen/eigvals). That commit touchesProject.toml,src/dual.jland tests, but notapiutils.jl. GitHub reports the PR as mergeable, but CI has not run on the merged tree.dense_seedableandSeededDualare new but unexported, so this is patch-level.test/GPUArraysTest.jl, which runs JLArrays withallowscalar(false), and adds JLArrays (MIT) as a test-only dependency with compat0.1, 0.2. Only 0.2 appears to have been exercised.Base.has_offset_axes, a Base internal that isn't marked public. The file already usesBase._unsetindex!, so there is precedent, but it is still a non-public dependency. No GPUArraysCore internals are used.lengthequality and then use linear views. That matches whateachindexdid before for dense buffers.gradient!at n=10 in vector mode going from 141 ns to 158 ns. The body only says "no directional claim".xthat is a non-contiguous wrapper (e.g.transpose(::CuArray)) goes throughview(x, linear_range)intomap!. This case is untested.Base.has_offset_axesis non-public; 1 commit behind master with no CI on the merged tree; no real CUDA run; a small vector-mode regression mentioned only in a comment; attribution lacks the model ID)Links
eigenandeigvals#757): a3c0f4fDenseArray/map!suggestion: Seed dense arrays of isbits duals without scalar indexing (fixes GPU jacobians) #816 (comment)gamma_incflake analysis: Seed dense arrays of isbits duals without scalar indexing (fixes GPU jacobians) #816 (comment)🤖 Risk assessment posted by an AI agent (fleet master) — harness: Devin CLI (local, Mac) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium; dispatched by Claude Code head, model claude-opus-5-5[1m]
Conversation: local Claude Code session 3cd6500a-1f81-46b5-ac0b-c466e15b6a53 on Chris's Mac (session ID, no URL)