Specialize split dual AD path on sparse matrix in ForwardDiff extension#1048
Conversation
…sion SparseArrays is a stdlib and has essentially zero installation cost. Making it a direct dependency (rather than a weakdep co-trigger) lets LinearSolveForwardDiffExt use `using SparseArrays` without SparseArrays needing to be loaded by the caller, so the ForwardDiff extension loads correctly for users who do not otherwise use SparseArrays.
…xCSC The generic update_partials_list! iterates over all m*n matrix elements and calls sparse_matrix[i,j] (a binary search into the column pointers) for each one. For a typical ODE Jacobian with low fill (e.g. 304 nnz out of 47*47=2209 entries) this is ~7x more work than necessary and uses O(log nnz) lookup per entry. Add SparseMatrixCSC specializations that iterate directly over the nzval array: - update_partials_list! reads nzval once and scatters each partial into the corresponding nzval of the pre-allocated output sparse matrices. - partials_to_list allocates output SparseMatrixCSC objects sharing the same sparsity pattern (copied colptr/rowval) with only the nzval varying per partial index. Both methods skip all structural zeros and avoid any index search, reducing the per-call cost from O(m*n * log(nnz)) to O(nnz).
…atrixCSC If the ODE Jacobian sparsity changes between factorization events, the cached list of partial matrices has different nnz than the new partial matrix, causing a DimensionMismatch in eachindex. Rebuild the cached matrices in-place when nnz changes.
At ODE init time the Jacobian may be structurally all-zeros, so nonzeros(partial_matrix) is empty. Use ForwardDiff.npartials/valtype on the element type instead of first(nz) to determine p and V.
|
Example timings from my package: using SymBoltz, ForwardDiff
M = ΛCDM()
p= parameters_Planck18(M)
prob = CosmologyProblem(M, p; jac = true, sparse = true) # stiff ODE with 82 unknowns and sparse+analytical Jacobian
vary = [M.c.Ω₀, M.b.Ω₀]
probgen = parameter_updater(prob, vary)
k = logrange(1e0, 1e3; length = 100)
P(θ) = spectrum_matter(probgen(θ), k; thread = false) # depends on solution of 100x of the above ODEs
θ₀ = [p[par] for par in vary]
@time ∂P_∂θ = ForwardDiff.jacobian(P, θ₀)
@profview ∂P_∂θ = ForwardDiff.jacobian(P, θ₀)Before this PR (note the big After this PR (note the small |
hersle
changed the title
|
Will fix formating. Is it OK to move SparseArrays to deps? That is the cleanest solution I could think of. Or do you see another way? |
|
Which linear solver? It shouldn't make a difference on PureKLU, and that would likely be pretty fast in a lot cases like this. |
|
Can do an extension that is SparseArrays+ForwardDiff for these. |
Using
Won't that end up duplicating a lot of the machinery in |
|
Oh wait, you want to do it with #1023 |
|
Oh yes, I was confused why I still saw that PR around. So wait, with the forwarddiff opt out you still don't see an advantage to PureKLU with forwarddiff? That's wild because it's a huge difference. |
|
No, I don't see any improvement with it. I tried the same code as above, now with 
This is after warm-up. Seems like its doing a lot of allocations, that could perhaps explain parts of it. |
|
The A matrix is dual? |
|
Yes it should be, its the |
|
That's really odd. |
|
Yeah I was also expecting it to perform well. Here is 
|
2 participants
Checklist
contributor guidelines, in particular the SciML Style Guide and
COLPRAC.
Additional context
In my package I need to forward-differentiate stiff ODEs with sparse Jacobians. Before this PR, the generic dual AD path through linear solves in the ODE has been a big bottleneck, as it does not specialize properly on the sparse matrix, but iterates through all dense$m \times n$ entries (including zeros) and goes through slow binary searches to look up the sparse matrix for each $(i, j)$ .
With this PR, I try to specialize the expensive
partials_to_listandupdate_partials_listfunctions for sparse matrices in order to remove this bottleneck.