From ee8884fe5938281e275d155e402db728ed319f99 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Wed, 17 Jun 2026 20:55:17 +0200 Subject: [PATCH 1/5] Make partials list dispatch signatures consistent --- ext/LinearSolveForwardDiffExt.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index b7925bb23..8e00d60ee 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -636,10 +636,9 @@ function update_partials_list!(partial_matrix::AbstractVector{T}, list_cache) wh return list_cache end -function update_partials_list!(partial_matrix, list_cache) +function update_partials_list!(partial_matrix::AbstractMatrix{T}, list_cache) where {T} p = length(first(partial_matrix)) m, n = size(partial_matrix) - for k in 1:p for i in 1:m for j in 1:n @@ -655,7 +654,7 @@ function partials_to_list(partial_matrix::AbstractVector{T}) where {T} return [[partial[i] for partial in partial_matrix] for i in p] end -function partials_to_list(partial_matrix) +function partials_to_list(partial_matrix::AbstractMatrix{T}) where {T} p = length(first(partial_matrix)) m, n = size(partial_matrix) res_list = fill(zeros(typeof(partial_matrix[1, 1][1]), m, n), p) From 1dc68ece519907aaf4595dff88f5ed35f3b86636 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Tue, 16 Jun 2026 17:46:23 +0200 Subject: [PATCH 2/5] Move SparseArrays to deps so it is available in the ForwardDiff extension 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. --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index c6a144b32..a23ab45ef 100644 --- a/Project.toml +++ b/Project.toml @@ -26,6 +26,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1" SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseColumnPivotedQR = "a57abbd0-fea5-4d57-96be-5e525945e8e4" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" @@ -62,7 +63,6 @@ PartitionedSolvers = "11b65f7f-80ac-401b-9ef2-3db765482d62" PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" STRUMPACK_jll = "86fbd0b9-476f-557c-b766-62c724b42d8c" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1" Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac" SpecializingFactorizations = "fa08b7a1-13d3-4faf-875d-5cbc1520e3f3" From 0abfd72cb0ae29d8be5044022e4f689344867675 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Tue, 16 Jun 2026 17:37:37 +0200 Subject: [PATCH 3/5] Specialize update_partials_list! and partials_to_list for SparseMatrixCSC 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). --- ext/LinearSolveForwardDiffExt.jl | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 8e00d60ee..b73901339 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -4,6 +4,7 @@ using LinearSolve using LinearSolve: SciMLLinearSolveAlgorithm, __init, LinearVerbosity, DefaultLinearSolver, DefaultAlgorithmChoice, defaultalg, reinit! using LinearAlgebra +using SparseArrays using ForwardDiff using ForwardDiff: Dual, Partials using SciMLBase @@ -670,4 +671,26 @@ function partials_to_list(partial_matrix::AbstractMatrix{T}) where {T} return res_list end +# Specializations for sparse matrices + +function partials_to_list(partial_matrix::SparseMatrixCSC) + nz = nonzeros(partial_matrix) + p = length(first(nz)) + V = typeof(first(nz)[1]) + m, n = size(partial_matrix) + return [SparseMatrixCSC(m, n, copy(partial_matrix.colptr), copy(partial_matrix.rowval), + V[nz[i][k] for i in eachindex(nz)]) for k in 1:p] +end + +function update_partials_list!(partial_matrix::SparseMatrixCSC, list_cache) + nz = nonzeros(partial_matrix) + for k in eachindex(list_cache) + nz_k = nonzeros(list_cache[k]) + @inbounds for i in eachindex(nz, nz_k) + nz_k[i] = nz[i][k] + end + end + return list_cache +end + end From a999104f71981bf83942664792e0b84f11829fa3 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Wed, 17 Jun 2026 11:46:01 +0200 Subject: [PATCH 4/5] Handle changing sparsity pattern in update_partials_list! for SparseMatrixCSC 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. --- ext/LinearSolveForwardDiffExt.jl | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index b73901339..dc1ad2a9b 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -684,10 +684,14 @@ end function update_partials_list!(partial_matrix::SparseMatrixCSC, list_cache) nz = nonzeros(partial_matrix) - for k in eachindex(list_cache) - nz_k = nonzeros(list_cache[k]) - @inbounds for i in eachindex(nz, nz_k) - nz_k[i] = nz[i][k] + if length(nz) != length(nonzeros(first(list_cache))) # TODO: more precise? + list_cache .= partials_to_list(partial_matrix) # sparsity pattern changed + else + for k in eachindex(list_cache) + nz_k = nonzeros(list_cache[k]) + @inbounds for i in eachindex(nz, nz_k) + nz_k[i] = nz[i][k] + end end end return list_cache From 0492ab1333486a43ccc8e9b66ab53c6251603d88 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Wed, 17 Jun 2026 11:48:21 +0200 Subject: [PATCH 5/5] Handle SparseMatrixCSC with empty nonzeros in partials_to_list 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. --- ext/LinearSolveForwardDiffExt.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index dc1ad2a9b..d21e86fcc 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -675,9 +675,10 @@ end function partials_to_list(partial_matrix::SparseMatrixCSC) nz = nonzeros(partial_matrix) - p = length(first(nz)) - V = typeof(first(nz)[1]) m, n = size(partial_matrix) + T = eltype(partial_matrix) + p = ForwardDiff.npartials(T) + V = ForwardDiff.valtype(T) # use type for concrete array below in empty-nz case (e.g. all-zero Jacobian at init) return [SparseMatrixCSC(m, n, copy(partial_matrix.colptr), copy(partial_matrix.rowval), V[nz[i][k] for i in eachindex(nz)]) for k in 1:p] end