Opt PureKLU and RFLU out of split dual AD path (rebase of #1023) - #1041
Merged
ChrisRackauckas merged 5 commits intoJun 14, 2026
Merged
Conversation
Two fixes for the direct dual opt-out path:
1. PureKLUFactorization's sparse init_cacheval methods were restricted to
Float64/ComplexF64 eltypes, so Dual-eltype sparse matrices fell through
to the nothing fallback and solve! crashed on cacheval.nzval. PureKLU's
kernels are pure Julia and generic over Union{Real, Complex}, so widen
the specializations accordingly.
2. The direct dual path built its inner problem from dual_A/dual_b as
stored, but when only A carries duals b is a plain array, and __init
takes the solution eltype from b, so the dual solution could not be
stored (MethodError: no method matching Float64(::Dual)). This also
affected the pre-existing GenericLUFactorization opt-out. Promote b to
the cache's dual type in _solve_direct_dual!, and only take the direct
path when A itself carries duals: with duals just in b, the split path
(one primal factorization + partials back-solves) is strictly cheaper
and works for all algorithms.
Adds test coverage for the duals-only-in-A and duals-only-in-b cases for
GenericLUFactorization, RFLUFactorization, and PureKLUFactorization.
Co-authored-by: Chris Rackauckas (Claude) <accounts@chrisrackauckas.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
The original PR widened the `{T, Int64}` / `{T, Int32}` PureKLU `init_cacheval`
methods from `Union{Float64, ComplexF64}` to `Union{Real, Complex}` to let
ForwardDiff duals dispatch on the direct dual solve path. After rebasing onto
main, that is unnecessary: SciML#1037 ("Default to PureKLU for generic-eltype sparse
LU") added a catch-all `where {T <: Number, Ti <: Integer}` method that already
builds the correct empty cache for any number eltype (duals included). The
widened specializations only duplicated it, so revert that file to main.
Verified: test/core/forwarddiff_overloads.jl (incl. the PureKLU sparse-dual
cases) passes with this file unchanged from main.
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
ChrisRackauckas
marked this pull request as ready for review
June 14, 2026 09:06
This was referenced Jun 18, 2026
ChrisRackauckas
added a commit
that referenced
this pull request
Jun 21, 2026
PR #1041 added RFLUFactorization to the direct dual solve path, taken whenever A carries duals. For ForwardDiff over an ODE solve (e.g. Rodas5P), the Rosenbrock W matrix carries duals, so every Newton solve factorized W in Dual arithmetic via RecursiveFactorization. RFLU's Float64 factorization is BLAS/SIMD-grade (cache-blocked, vectorized); routing the Dual problem through it falls back to generic scalar dual arithmetic and loses that speedup entirely. Measured on the issue's case: ~40x slower (0.07s -> 2.9s) and ~13x more allocation (15.6 MiB -> 199.9 MiB), with identical derivatives. The split primal/partials path already handles duals-in-A correctly: it factorizes the primal W once and reuses that factorization across the partial back-solves, which is what 3.85.1 did. Drop RFLU from _use_direct_dual_solve so it stays on that path. GenericLU/SpecializedLU/ SpecializedQR (cheap in dual arithmetic) and PureKLU are unaffected. Adds a regression test asserting RFLU is not routed to the direct path. Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChrisRackauckas
added a commit
that referenced
this pull request
Jun 30, 2026
… 1.1 mixed ldiv!) (#1069) * Add mixed-type ldiv! for primal KLU factor against a Dual RHS A factorization is a linear operator, so `K \ Dual(v, p₁…p_N) = Dual(K\v, K\p₁ … K\p_N)`. PureKLU stores its L/U in the factor element type, so a `KLUFactorization{<:AbstractFloat}` can backsolve a Dual right-hand side by solving the primal value and each partial column in one multi-RHS `klu_solve!`, then repacking — without ever promoting `A` to a Dual. This keeps the factorization in fast Float64 arithmetic and lets the duals ride only through the back-substitution, which is the single-solve analogue of the split path's per-partial back-solves. It makes `solve(LinearProblem(A_float64, b_dual), PureKLU)` work natively (previously a MethodError on `ldiv!(::KLUFactorization{Float64}, ::Vector{<:Dual})`). The factor eltype is constrained to `AbstractFloat` so it never collides with PureKLU's own `ldiv!(::KLUFactorization{Tv}, ::VecOrMat{Tv})` on the dual-factor path — a `Dual` is `<: Real` but not `<: AbstractFloat`, so the dual-factor case stays with PureKLU's method. Adds a correctness test (value + every partial column matches the dense Dual solve, factor stays Float64) for chunk sizes 1/2/3. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Route duals-in-b PureKLU solves to a native LinearCache With the mixed-type `ldiv!` (primal KLU factor against a Dual RHS) in place, a `DualBLinearProblem` (duals only in b, A primal) no longer needs the split DualLinearCache machinery: it can be solved on a plain `LinearCache` whose factorization stays in Float64, with the duals carried through the back-solve. Add an `init(::DualBLinearProblem, ::PureKLUFactorization)` opt-out that returns `__init(prob, alg)` directly. This is type-stable — dispatch is purely on the problem subtype (b-dual / A-plain) and the alg type, so `init` always returns a `LinearCache` for this method (verified with `@inferred`), unlike the reverted value-based opt-out in #1041-era code. It is also simpler than the split path and gets correct factorization reuse across b-only `reinit!`s for free, which the split path's `reinit!` does not (it unconditionally marks the inner cache fresh, forcing a refactorization even when only b changed). Scope: only PureKLU is routed, since it is the only sparse factorization with the mixed `ldiv!`. A-dual and mixed A/b problems are untouched and still take the split path. Adds routing + type-stability + correctness tests. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Move the primal-factor/Dual-RHS ldiv! to PureKLU; depend on PureKLU 1.1 The mixed-type `ldiv!` (primal KLU factor backsolving a Dual RHS) belongs in PureKLU, which owns `KLUFactorization` — it is the analogue of PureKLU's existing real-factor / Complex-RHS `ldiv!`, avoids type piracy here, and there it is non-allocating (the real `n × (N+1)` RHS reuses a buffer on the factorization). See SciML/PureKLU.jl#67. This drops the stop-gap `ldiv!` (and the `PureKLU` import) from this extension and bumps the PureKLU [compat] to 1.1, which provides it. The duals-in-b routing (`init(::DualBLinearProblem, ::PureKLUFactorization)` -> native `LinearCache`) is unchanged and now relies on PureKLU's `ldiv!`. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Pull PureKLU 1.1.0 from PR #67 via [sources] so #1069 CI resolves The duals-in-b routing needs the primal-factor/Dual-RHS `ldiv!` from PureKLU 1.1.0 (SciML/PureKLU.jl#67), which is not yet registered. Add a temporary `[sources]` entry pointing at the PR branch so `PureKLU = "1.1"` resolves in CI (Julia >= 1.11; ignored on older Julia). To be removed once PureKLU 1.1.0 is in the General registry. Verified locally: a fresh resolve with this `[sources]` pulls PureKLU v1.1.0. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Drop temporary [sources] now that PureKLU 1.1.0 is registered PureKLU 1.1.0 (with the primal-factor/Dual-RHS `ldiv!`, SciML/PureKLU.jl#67) is in the General registry, so the `PureKLU = "1.1"` compat resolves directly. Remove the temporary `[sources]` override that pointed at the PR branch. Verified: forwarddiff_overloads.jl passes end-to-end against the registered PureKLU 1.1.0. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> * Apply suggestion from @ChrisRackauckas --------- Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rebases and finalizes #1023 (branch
ChrisRackauckas-patch-1) onto the currentmain. The original PR went stale (mergeable_state: dirty); I couldn't push to the SciML-owned branch, so this is the rebased version from a fork. It carries the same commits (authored by @ChrisRackauckas), now conflict-free onmain. Supersedes #1023 — close that one in favor of this.What it does (unchanged intent)
Opt
PureKLUFactorizationandRFLUFactorizationout of the split primal/partials dual-AD path so they solve the Dual problem directly, and only take the direct path whenAitself carries duals (with duals only inb, the split path — one primal factorization + partials back-solves — is cheaper than factorizing in dual arithmetic). Also promotes a plainbto the dual eltype on the direct path.Net change is now just
LinearSolveForwardDiffExt.jl+ the testThe original PR also widened two PureKLU
init_cachevalmethods inLinearSolveSparseArraysExt.jlfromUnion{Float64, ComplexF64}toUnion{Real, Complex}. That is dropped here — after rebasing ontomain, it's redundant: #1037 ("Default to PureKLU for generic-eltype sparse LU") added a catch-allwhere {T <: Number, Ti <: Integer}method that already builds the correct empty cache for any number eltype, duals included. The widened specializations only duplicated it, so that file is left identical tomain. (The last commit reverts it, with the reasoning; verified the PureKLU sparse-dual tests still pass.)Conflicts resolved during rebase
mainindependently evolved both touched spots:ext/LinearSolveForwardDiffExt.jl—_use_direct_dual_solvehad gainedSpecializedLUFactorization/SpecializedQRFactorizationonmain. Resolved as the union:GenericLU || SpecializedLU || SpecializedQR || PureKLU || RFLU.test/core/forwarddiff_overloads.jl—mainswitched the direct-path comparison to a robustdual_isapproxhelper (plain≈over Dual vectors NaNs when a primal diff is exactly zero). Kept that helper and theSpecialized*+ least-squares tests, and re-expressed the PR'sGenericLU/RFLUchecks throughdual_isapprox(they take the same direct path, so the same ulp/NaN concern applies). Unioned theusing SpecializingFactorizations/using RecursiveFactorizationimports.Local verification
Julia 1.12 / Linux.
Runic --checkclean.test/core/forwarddiff_overloads.jlruns green end-to-end (all testsets, including the mergeddual_isapproxchecks forSpecializedLU/SpecializedQR/GenericLU/RFLU, the least-squares case, the plain-A/plain-bopt-out cases, and thePureKLUsparse-dual tests) — both with and without the now-dropped ext change.Please ignore until reviewed by @ChrisRackauckas.
🤖 Generated with Claude Code