From 63a19cf677f818f322e0e919c6e00253a2b57c35 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 28 Jun 2026 10:09:45 -0400 Subject: [PATCH] Keep PureKLUFactorization off the direct dual AD path (fixes #1064) PR #1041 added PureKLUFactorization to the direct dual solve path, taken whenever A carries duals. For ForwardDiff over an implicit ODE solve (e.g. Rodas5P), the Rosenbrock W matrix carries duals, so every solve factorized W in dual arithmetic via PureKLU's pure-Julia KLU. This is the same problem that #1052/#1057 fixed for RFLUFactorization, which was opted out at the same time; PureKLU was left on the direct path ("PureKLU are unaffected"), but #1064 shows it is affected identically. The split primal/partials path factorizes the primal W once in fast Float64 arithmetic and reuses that factorization across the partial back-solves, instead of re-running the whole numeric factorization in (2N+1)-wide scalar dual arithmetic. On the issue's reproducer (tridiagonal heat equation, N=20, Rodas5P, single derivative) the direct path was ~2.4x slower (3.79ms -> 1.56ms) with ~35x the allocation (7.93 MiB -> 227 KiB), and the split path matches the trusted C-backed KLUFactorization (1.49ms / 222 KiB), with identical derivatives. Drop PureKLU from _use_direct_dual_solve so it stays on the split path, and add a regression test asserting it is not routed to the direct path. Note: the split path is the better default for the sparse-Jacobian / AD-through-implicit-ODE workload PureKLU+ForwardDiff is used in. For a single solve of a very-low-fill matrix with a large dual chunk the direct path can edge ahead, but that is not the regime this hits; a separate inefficiency in _solve_direct_dual! (it re-inits a fresh cache on every solve, defeating KLU symbolic reuse) is the real source of the large allocation and is left for a follow-up. Co-Authored-By: Chris Rackauckas --- ext/LinearSolveForwardDiffExt.jl | 20 +++++++++++++------- test/Core/forwarddiff_overloads.jl | 10 ++++++++++ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 4acc7e9ad..8fcfe8e2d 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -393,15 +393,21 @@ end # Check if the algorithm should use the direct dual solve path # (algorithms that can work directly with Dual numbers without the primal/partials separation) function _use_direct_dual_solve(alg) - # NOTE: RFLUFactorization is intentionally *not* on the direct path. Even when - # A carries duals, its fast Float64 factorization is BLAS/SIMD-grade, and routing - # the Dual problem through it falls back to generic scalar dual arithmetic, losing - # that speedup entirely (~40x slower, see issue #1052). The split path keeps the - # fast primal factorization and reuses it across the partial back-solves. + # NOTE: RFLUFactorization and PureKLUFactorization are intentionally *not* on the + # direct path. The split path factorizes the *primal* A once (in fast Float64 + # arithmetic) and reuses that factorization across the partial back-solves, instead + # of re-running the whole numeric factorization in (2N+1)-wide scalar dual + # arithmetic. For RFLU this is decisive because its Float64 factorization is + # BLAS/SIMD-grade (~40x slower on the direct path, see #1052). For PureKLU it wins + # whenever the factorization carries non-trivial fill (general sparse Jacobians, + # PDE stencils) and for the small-chunk AD-through-implicit-ODE workload these are + # actually used in — see #1064, where the direct path was ~2.4x slower with ~35x + # the allocation through a Rosenbrock solve, matching the trusted C-KLU split path + # once routed here. (For a single solve of a very-low-fill matrix with a large dual + # chunk the direct path can edge ahead, but that is not the regime PureKLU+AD hits.) return alg isa GenericLUFactorization || alg isa LinearSolve.SpecializedLUFactorization || - alg isa LinearSolve.SpecializedQRFactorization || - alg isa LinearSolve.PureKLUFactorization + alg isa LinearSolve.SpecializedQRFactorization end function _use_direct_dual_solve(alg::DefaultLinearSolver) diff --git a/test/Core/forwarddiff_overloads.jl b/test/Core/forwarddiff_overloads.jl index 978a4faa3..3c1339f61 100644 --- a/test/Core/forwarddiff_overloads.jl +++ b/test/Core/forwarddiff_overloads.jl @@ -280,6 +280,16 @@ plain_A = ForwardDiff.value.(A) prob = LinearProblem(sparse(plain_A), b) @test ≈(solve(prob, PureKLUFactorization()), plain_A \ b, rtol = 1.0e-9) +# Regression test for #1064: PureKLUFactorization must stay on the split +# primal/partials path and NOT take the direct dual solve. Like RFLU (#1052), +# its fast Float64 (KLU) factorization is far cheaper than factorizing the Dual +# problem in generic scalar dual arithmetic; the direct path was much slower with +# ~14x more allocation through an implicit ODE solve. Guard the routing decision. +@testset "PureKLU stays off the direct dual path (#1064)" begin + ext = Base.get_extension(LinearSolve, :LinearSolveForwardDiffExt) + @test !ext._use_direct_dual_solve(PureKLUFactorization()) +end + A, b = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)]) prob = LinearProblem(sparse(A), sparse(b))