Route RFLU matrix right-hand sides through TriangularSolve (#1117)

RecursiveFactorization already uses TriangularSolve inside its own lu!,
but the ldiv! that consumes the factorization only does so for the
pivotless NotIPIV case (RecursiveFactorization/src/lu.jl); a pivoted LU -
which is what RFLUFactorization produces - falls back to LinearAlgebra,
i.e. BLAS trsv/trsm.

Fixing that inside RecursiveFactorization would be type piracy (it owns
neither ldiv! nor LU), so it belongs here: LinearSolve owns solve! for
its own algorithm type. Matrix right-hand sides now go through
TriangularSolve's blocked kernels.

Measured solve-only, factorization already computed, 1 BLAS thread:

  n     nrhs=4                nrhs=8
  64    6.31 -> 3.76 us       8.76 -> 5.40 us    (1.68x / 1.62x)
  128  21.11 -> 12.68 us     26.20 -> 17.22 us   (1.66x / 1.52x)
  256  86.43 -> 55.11 us     97.09 -> 67.87 us   (1.57x / 1.43x)
  500 277.11 -> 175.42 us   301.13 -> 225.77 us  (1.58x / 1.33x)

Vector right-hand sides keep the stdlib path: TriangularSolve has no
vector kernel, and reshaping to n x 1 measured 1.09x at n=128 but 0.88x
by n=256. A single-column matrix is the same case in disguise, so it is
guarded explicitly - with the guard, nrhs=1 and vector results are
bit-identical to the stdlib path (agreement exactly 0.0), while
multi-RHS keeps 1.5-1.7x.

TriangularSolve is added as a weakdep purely to trigger the extension;
RecursiveFactorization depends on it, so loading RecursiveFactorization
alone still activates everything and no dependency graph changes.

New multi-RHS RFLU testset covers the direct and cache-reuse paths.

Co-authored-by: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

Project.toml

@@ -65,6 +65,7 @@ PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 PartitionedSolvers = "11b65f7f-80ac-401b-9ef2-3db765482d62"
 PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b"
 RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
+TriangularSolve = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf"
 STRUMPACK_jll = "86fbd0b9-476f-557c-b766-62c724b42d8c"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac"
@@ -106,7 +107,7 @@ LinearSolvePartitionedSolversExt = ["PartitionedArrays", "PartitionedSolvers"]
 LinearSolveParUExt = ["ParU_jll", "SparseArrays"]
 LinearSolvePureUMFPACKExt = ["PureUMFPACK", "SparseArrays"]
 LinearSolvePardisoExt = ["Pardiso", "SparseArrays"]
-LinearSolveRecursiveFactorizationExt = "RecursiveFactorization"
+LinearSolveRecursiveFactorizationExt = ["RecursiveFactorization", "TriangularSolve"]
 LinearSolveSuperLUDISTExt = ["SparseArrays", "SuperLUDIST"]
 LinearSolveSTRUMPACKExt = ["SparseArrays", "STRUMPACK_jll"]
 LinearSolveSparseArraysExt = "SparseArrays"
@@ -174,6 +175,7 @@ PureUMFPACK = "0.1"
 Random = "1.10"
 RecursiveArrayTools = "3.37, 4"
 RecursiveFactorization = "0.2.26"
+TriangularSolve = "0.2"
 Reexport = "1.2.2"
 STRUMPACK_jll = "8"
 SafeTestsets = "0.1"

ext/LinearSolveRecursiveFactorizationExt.jl

@@ -4,6 +4,7 @@ using LinearSolve: LinearSolve, userecursivefactorization, LinearCache, @get_cac
     RFLUFactorization, ButterflyFactorization, RF32MixedLUFactorization,
     LinearVerbosity
 using LinearSolve.LinearAlgebra, LinearSolve.ArrayInterface, RecursiveFactorization
+using TriangularSolve: TriangularSolve
 using SciMLBase: SciMLBase, ReturnCode
 using SciMLLogging: @SciMLMessage
 
@@ -31,10 +32,59 @@ function SciMLBase.solve!(
 
         cache.isfresh = false
     end
-    y = ldiv!(cache.u, LinearSolve.@get_cacheval(cache, :RFLUFactorization)[1], cache.b)
+    y = _rf_ldiv!(
+        cache.u, LinearSolve.@get_cacheval(cache, :RFLUFactorization)[1], cache.b, Val(T)
+    )
     return SciMLBase.build_linear_solution(alg, y, nothing, nothing; retcode = ReturnCode.Success)
 end
 
+# Apply an RF factorization to a right-hand side.
+#
+# `RecursiveFactorization` already routes its own `lu!` through TriangularSolve,
+# but the `ldiv!` that consumes the factorization only does so for the pivotless
+# `NotIPIV` case (RecursiveFactorization/src/lu.jl); a pivoted `LU` falls back to
+# LinearAlgebra, i.e. BLAS `trsv`/`trsm`.  For a matrix right-hand side that
+# leaves a consistent ~1.6x on the table at every size we measured, because
+# TriangularSolve's blocked kernels beat `trsm` here:
+#
+#   n     BLAS trsm   TriangularSolve
+#   32     2.19 us      1.37 us  (1.60x)
+#   64     6.27 us      3.76 us  (1.67x)
+#   128   21.02 us     12.72 us  (1.65x)
+#   256   86.08 us     54.07 us  (1.59x)
+#   500  276.20 us    175.81 us  (1.58x)
+#
+# For a single vector right-hand side TriangularSolve has no advantage (it has
+# no vector kernel, and reshaping to n x 1 measured 1.09x at n=128 but 0.88x by
+# n=256), so vectors keep the stdlib path.
+@inline function _rf_ldiv!(
+        u::AbstractVector, fact::LinearAlgebra.LU, b::AbstractVector, ::Val
+    )
+    return ldiv!(u, fact, b)
+end
+
+function _rf_ldiv!(
+        U::AbstractMatrix{T}, fact::LinearAlgebra.LU{T, <:StridedMatrix{T}},
+        B::AbstractMatrix{T}, ::Val{Thread}
+    ) where {T <: LinearAlgebra.BlasFloat, Thread}
+    # A single column is the vector case in disguise: measured 0.87-0.90x at
+    # n >= 256, so it keeps the stdlib path.
+    size(B, 2) == 1 && return ldiv!(U, fact, B)
+    U === B || copyto!(U, B)
+    LinearAlgebra._ipiv_rows!(fact, 1:length(fact.ipiv), U)
+    F = fact.factors
+    TriangularSolve.ldiv!(UnitLowerTriangular(F), U, Val(Thread))
+    TriangularSolve.ldiv!(UpperTriangular(F), U, Val(Thread))
+    return U
+end
+
+# Non-strided or non-BLAS element types keep the stdlib path.
+@inline function _rf_ldiv!(
+        U::AbstractMatrix, fact::LinearAlgebra.LU, B::AbstractMatrix, ::Val
+    )
+    return ldiv!(U, fact, B)
+end
+
 # Mixed precision RecursiveFactorization implementation
 
 const PREALLOCATED_RF32_LU = begin

test/Core/basictests.jl

@@ -211,6 +211,24 @@ end
         @test X * solve!(cache) ≈ b1
     end
 
+    @testset "RFLUFactorization multi-RHS" begin
+        # The extension routes matrix right-hand sides through TriangularSolve;
+        # results must match the stdlib path exactly in behaviour.
+        Random.seed!(7)
+        for n in (8, 40), nrhs in (1, 3)
+            Ar = rand(n, n) + n * I
+            Br = rand(n, nrhs)
+            sol = solve(LinearProblem(copy(Ar), copy(Br)), RFLUFactorization())
+            @test sol.u ≈ Ar \ Br
+            # cache reuse path
+            cache = SciMLBase.init(LinearProblem(copy(Ar), copy(Br)), RFLUFactorization())
+            @test solve!(cache).u ≈ Ar \ Br
+            A2 = Ar + I
+            cache.A = copy(A2)
+            @test solve!(cache).u ≈ A2 \ Br
+        end
+    end
+
     @testset "SupernodalLU Factorization" begin
         A1 = sparse(A / 1)
         b1 = rand(n)