From 3b6916b1d4b34fc6d165bae9804795ed1a7b3041 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 21 Jun 2026 06:18:11 -0400 Subject: [PATCH] Keep RFLUFactorization off the direct dual AD path (fixes #1052) 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: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) --- Project.toml | 2 +- ext/LinearSolveForwardDiffExt.jl | 8 ++++++-- test/Core/forwarddiff_overloads.jl | 13 +++++++++++++ 3 files changed, 20 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 3333dd866..5f8f811ee 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "LinearSolve" uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" -version = "3.86.0" +version = "3.86.1" authors = ["SciML"] [deps] diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index b15924785..97d03d7bd 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -393,11 +393,15 @@ 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. return alg isa GenericLUFactorization || alg isa LinearSolve.SpecializedLUFactorization || alg isa LinearSolve.SpecializedQRFactorization || - alg isa LinearSolve.PureKLUFactorization || - alg isa LinearSolve.RFLUFactorization + alg isa LinearSolve.PureKLUFactorization end function _use_direct_dual_solve(alg::DefaultLinearSolver) diff --git a/test/Core/forwarddiff_overloads.jl b/test/Core/forwarddiff_overloads.jl index 1a93d2a54..373f60e20 100644 --- a/test/Core/forwarddiff_overloads.jl +++ b/test/Core/forwarddiff_overloads.jl @@ -67,6 +67,19 @@ prob = LinearProblem(plain_A, b) @test ≈(solve(prob, GenericLUFactorization()), plain_A \ b, rtol = 1.0e-9) @test ≈(solve(prob, RFLUFactorization()), plain_A \ b, rtol = 1.0e-9) +# Regression test for #1052: RFLUFactorization must stay on the split +# primal/partials path and NOT take the direct dual solve. Its fast Float64 +# factorization is BLAS/SIMD-grade; routing the Dual problem through it falls +# back to generic scalar dual arithmetic (~40x slower). Guard the routing +# decision directly so RFLU is never re-added to the direct path. +@testset "RFLU stays off the direct dual path (#1052)" begin + ext = Base.get_extension(LinearSolve, :LinearSolveForwardDiffExt) + @test !ext._use_direct_dual_solve(RFLUFactorization()) + # Sanity: the genuinely-cheap-in-dual algorithms are still on the direct path. + @test ext._use_direct_dual_solve(GenericLUFactorization()) + @test ext._use_direct_dual_solve(SpecializedLUFactorization()) +end + # Overload Dense A, b = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)])