Batched (matrix) right-hand sides: LinearProblem(A, B) ≡ A \ B (v4.0.0)

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearSolve"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
-version = "3.87.0"
+version = "4.0.0"
 authors = ["SciML"]
 
 [deps]
@@ -110,7 +110,10 @@ LinearSolveSpecializingFactorizationsExt = "SpecializingFactorizations"
 [compat]
 AMD = "0.5"
 AMDGPU = "1.2, 2"
-AlgebraicMultigrid = "1, 2"
+# 0.5 allowed so the test target resolves against LinearSolve 4 until
+# AlgebraicMultigrid widens its own LinearSolve compat (its 1.x/2.x cap at 3);
+# the ext API is identical and the AMG testset passes on 0.5.1.
+AlgebraicMultigrid = "0.5, 1, 2"
 ArrayInterface = "7.19"
 BandedMatrices = "1.8"
 BlockDiagonals = "0.2"

docs/Project.toml

@@ -6,6 +6,10 @@ SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 
 [compat]
 Documenter = "1"
-LinearSolve = "3"
+LinearSolve = "4"
 LinearSolveAutotune = "1.1"
 SciMLOperators = "1"
+
+[sources]
+LinearSolve = {path = ".."}
+LinearSolveAutotune = {path = "../lib/LinearSolveAutotune"}

docs/src/release_notes.md

@@ -1,5 +1,14 @@
 # Release Notes
 
+## v4.0
+
+  - Batched (matrix) right-hand sides are now supported: `solve(LinearProblem(A, B))` with
+    `B::AbstractMatrix` computes the equivalent of `A \ B`, factorizing `A` once and
+    returning `sol.u` as a `size(A, 2) × size(B, 2)` matrix. This is a breaking change:
+    previously a matrix `b` initialized a vector-shaped `u` and generally errored downstream.
+    Batched right-hand sides are supported by the factorization-based algorithms; iterative
+    (Krylov) methods throw an informative `ArgumentError` for matrix `b`.
+
 ## Upcoming Changes
 
   - `CudaOffloadFactorization` has been split into two algorithms:

docs/src/tutorials/linear.md

@@ -36,6 +36,26 @@ pass in an algorithm struct and all wrapped linear solvers are immediately
 available as tweaks to the general algorithm. For more information on the
 available solvers, see [the solvers page](@ref linearsystemsolvers)
 
+## Batched Right-Hand Sides
+
+A matrix right-hand side `B` solves all of its columns against the same `A` at
+once, just like `A \ B`: `solve(LinearProblem(A, B))` factorizes `A` a single
+time and returns `sol.u` as the `size(A, 2) × size(B, 2)` matrix satisfying
+`A * sol.u ≈ B`.
+
+```@example linsys1
+B = rand(4, 2)
+prob = LS.LinearProblem(A, B)
+sol = LS.solve(prob)
+sol.u ≈ A \ B
+```
+
+This works with the caching interface as well (`cache.b = B2; solve!(cache)`
+re-solves all columns against the cached factorization). Batched right-hand
+sides are only supported by factorization-based algorithms; iterative (Krylov)
+methods such as `LS.KrylovJL_GMRES()` accept only vector `b` and throw an
+informative error for matrix `B` — solve column-by-column instead in that case.
+
 ## Sparse and Structured Matrices
 
 There is no difference in the interface for LinearSolve.jl on sparse

ext/LinearSolveSparseArraysExt.jl

@@ -778,12 +778,24 @@ function SciMLBase.solve!(
         cache.isfresh = false
     end
     F = LinearSolve.@get_cacheval(cache, :SparseColumnPivotedQRFactorization)
-    y = ldiv!(cache.u, F, cache.b)
+    y = LinearSolve._ldiv!(cache.u, F, cache.b)
     return SciMLBase.build_linear_solution(
         alg, y, nothing, cache; retcode = ReturnCode.Success
     )
 end
 
+# SparseColumnPivotedQR's ldiv! only accepts vector right-hand sides; batched
+# (matrix) right-hand sides solve column-by-column against the one factorization.
+function LinearSolve._ldiv!(
+        x::AbstractMatrix,
+        F::SCPQR.SparseColumnPivotedQRFactorization, b::AbstractMatrix
+    )
+    for j in axes(b, 2)
+        ldiv!(view(x, :, j), F, view(b, :, j))
+    end
+    return x
+end
+
 # Build a column-pivoted sparse QR factorization for the default sparse-LU
 # singular fallback (`_do_sparse_qr_fallback` in src/default.jl).
 function LinearSolve.sparse_colpivqr_factorize(A)
@@ -874,8 +886,8 @@ end
             x .= A \ b
         end
         function LinearSolve._ldiv!(
-                x::AbstractVector,
-                A::SparseArrays.CHOLMOD.Factor, b::AbstractVector
+                x::AbstractVecOrMat,
+                A::SparseArrays.CHOLMOD.Factor, b::AbstractVecOrMat
             )
             x .= A \ b
         end
@@ -897,6 +909,15 @@ end
         end
     end
 
+    # SPQR has no in-place matrix (batched) ldiv! on any current Julia version,
+    # so route batched right-hand sides through the allocating `\`.
+    function LinearSolve._ldiv!(
+            x::AbstractMatrix,
+            A::SparseArrays.SPQR.QRSparse, b::AbstractMatrix
+        )
+        x .= A \ b
+    end
+
     function LinearSolve._ldiv!(
             ::SVector,
             A::Union{SparseArrays.CHOLMOD.Factor, SparseArrays.SPQR.QRSparse},
@@ -940,11 +961,11 @@ end
 # medium and very sparse / "less structure") favors the pure-Julia KLU-style
 # solvers: PureKLU for LU and SparseColumnPivotedQR for QR. `false` ("more
 # structure") favors the SuiteSparse solvers: UMFPACK for LU and SPQR for QR.
-# The `length(b) <= 1_000` branch is the "small enough" fast path; the second is
+# The `size(b, 1) <= 1_000` branch is the "small enough" fast path; the second is
 # the "medium and sufficiently sparse" rule.
 function LinearSolve.use_klulike_sparse_structure(A::AbstractSparseMatrixCSC, b)
-    return length(b) <= 1_000 ||
-        (length(b) <= 10_000 && length(nonzeros(A)) / length(A) < 2.0e-4)
+    return size(b, 1) <= 1_000 ||
+        (size(b, 1) <= 10_000 && length(nonzeros(A)) / length(A) < 2.0e-4)
 end
 
 @static if Base.USE_GPL_LIBS

lib/LinearSolveAutotune/Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolveAutotune"
 uuid = "67398393-80e8-4254-b7e4-1b9a36a3c5b6"
 authors = ["SciML"]
-version = "1.11.1"
+version = "1.12.0"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
@@ -44,7 +44,7 @@ FastLapackInterface = "2.0.4"
 GitHub = "5"
 LAPACK_jll = "3.12"
 LinearAlgebra = "1"
-LinearSolve = "3.39.2"
+LinearSolve = "4"
 MKL_jll = "2025.2.0"
 Metal = "1.5"
 OpenBLAS_jll = "0.3"

lib/LinearSolveAutotune/test/qa/Project.toml

@@ -13,7 +13,7 @@ LinearSolveAutotune = {path = "../.."}
 [compat]
 Aqua = "0.8"
 JET = "0.9, 0.10, 0.11"
-LinearSolve = "3"
+LinearSolve = "4"
 LinearSolveAutotune = "1"
 SciMLTesting = "1.6"
 Test = "1"

lib/LinearSolvePyAMG/Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolvePyAMG"
 uuid = "7a56c47d-7ab1-4e99-b0e3-2952e463d64a"
 authors = ["SciML"]
-version = "1.1.1"
+version = "1.2.0"
 
 [deps]
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
@@ -17,7 +17,7 @@ LinearSolve = {path = "../.."}
 [compat]
 CondaPkg = "0.2"
 LinearAlgebra = "1.10"
-LinearSolve = "3"
+LinearSolve = "4"
 Pkg = "1"
 PythonCall = "0.9"
 SafeTestsets = "0.1"

lib/LinearSolvePyAMG/test/qa/Project.toml

@@ -13,7 +13,7 @@ LinearSolvePyAMG = {path = "../.."}
 [compat]
 Aqua = "0.8"
 JET = "0.9, 0.10, 0.11"
-LinearSolve = "3"
+LinearSolve = "4"
 LinearSolvePyAMG = "1"
 SciMLTesting = "1.6"
 Test = "1"

src/appleaccelerate.jl

@@ -382,7 +382,11 @@ function SciMLBase.solve!(
     if m > n
         Bc = copy(cache.b)
         aa_getrs!('N', A.factors, A.ipiv, Bc; info)
-        copyto!(cache.u, 1, Bc, 1, n)
+        if cache.b isa AbstractMatrix
+            copyto!(cache.u, @view(Bc[1:n, :]))
+        else
+            copyto!(cache.u, 1, Bc, 1, n)
+        end
     else
         copyto!(cache.u, cache.b)
         aa_getrs!('N', A.factors, A.ipiv, cache.u; info)
@@ -468,7 +472,11 @@ function SciMLBase.solve!(
     if m > n
         aa_getrs!('N', A_lu.factors, A_lu.ipiv, b_32; info)
         # Convert back to original precision
-        cache.u[1:n] .= Torig.(b_32[1:n])
+        if cache.b isa AbstractMatrix
+            cache.u .= Torig.(@view(b_32[1:n, :]))
+        else
+            cache.u[1:n] .= Torig.(@view(b_32[1:n]))
+        end
     else
         copyto!(u_32, b_32)
         aa_getrs!('N', A_lu.factors, A_lu.ipiv, u_32; info)

src/common.jl

@@ -405,6 +405,9 @@ same element type as `b` and sized to match the number of columns in `A`.
 
 ## Returns
 A zero-initialized vector of size `(size(A, 2),)` with element type matching `b`.
+For a matrix (batched) right-hand side `b` of size `(size(A, 1), k)`, returns a
+zero-initialized matrix of size `(size(A, 2), k)` so that each column of `u0`
+corresponds to a column of `b`.
 
 ## Specializations
 - For static matrices (`SMatrix`): Returns a static vector (`SVector`)
@@ -415,7 +418,56 @@ function __init_u0_from_Ab(A, b)
     fill!(u0, false)
     return u0
 end
+function __init_u0_from_Ab(A, b::AbstractMatrix)
+    u0 = similar(b, size(A, 2), size(b, 2))
+    fill!(u0, false)
+    return u0
+end
 __init_u0_from_Ab(::SMatrix{S1, S2}, b) where {S1, S2} = zeros(SVector{S2, eltype(b)})
+function __init_u0_from_Ab(::SMatrix{S1, S2}, b::AbstractMatrix) where {S1, S2}
+    u0 = similar(b, S2, size(b, 2))
+    fill!(u0, false)
+    return u0
+end
+function __init_u0_from_Ab(
+        ::SMatrix{S1, S2}, ::SMatrix{S1b, S2b, Tb}
+    ) where {S1, S2, S1b, S2b, Tb}
+    return zeros(SMatrix{S2, S2b, Tb})
+end
+
+"""
+    _check_batched_rhs_support(alg, b)
+
+Throw an informative `ArgumentError` at `init` time when a matrix (batched)
+right-hand side `b` is used with an algorithm that only supports vector `b`
+(Krylov subspace / iterative methods). Factorization-based algorithms support
+matrix `b` and pass through the generic no-op fallback.
+"""
+_check_batched_rhs_support(alg, b) = nothing
+function _check_batched_rhs_support(alg::AbstractKrylovSubspaceMethod, b::AbstractMatrix)
+    throw(
+        ArgumentError(
+            "Batched (matrix) right-hand sides are only supported by factorization " *
+                "algorithms; $(nameof(typeof(alg))) supports only vector `b`. Solve " *
+                "column-by-column or use a factorization algorithm (e.g. `LUFactorization()`)."
+        )
+    )
+end
+function _check_batched_rhs_support(alg::DefaultLinearSolver, b::AbstractMatrix)
+    if alg.alg === DefaultAlgorithmChoice.KrylovJL_GMRES ||
+            alg.alg === DefaultAlgorithmChoice.KrylovJL_CRAIGMR ||
+            alg.alg === DefaultAlgorithmChoice.KrylovJL_LSMR
+        throw(
+            ArgumentError(
+                "Batched (matrix) right-hand sides are only supported by factorization " *
+                    "algorithms; the default algorithm selected the Krylov method " *
+                    "$(alg.alg) for this operator, which supports only vector `b`. " *
+                    "Solve column-by-column or use a factorization algorithm."
+            )
+        )
+    end
+    return nothing
+end
 
 function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm, args...; kwargs...)
     return __init(prob, alg, args...; kwargs...)
@@ -506,6 +558,8 @@ function __init(
         copy(b)
     end
 
+    _check_batched_rhs_support(alg, b)
+
     u0_ = u0 !== nothing ? u0 : __init_u0_from_Ab(A, b)
 
     # Guard against type mismatch for user-specified reltol/abstol

src/default.jl

@@ -354,11 +354,13 @@ function defaultalg(A, b, assump::OperatorAssumptions{Bool})
                 )
 
                 # Small matrix override - always use GenericLUFactorization for tiny problems
-                if length(b) <= 10
+                # `size(b, 1)` (not `length(b)`) so batched (matrix) right-hand
+                # sides don't inflate the apparent problem size.
+                if size(b, 1) <= 10
                     DefaultAlgorithmChoice.GenericLUFactorization
                 else
                     # Check if autotune preferences exist for larger matrices
-                    matrix_size = length(b)
+                    matrix_size = size(b, 1)
                     eltype_A = A === nothing ? Nothing : eltype(A)
                     tuned_alg = get_tuned_algorithm(eltype_A, eltype(b), matrix_size)
 
@@ -369,8 +371,8 @@ function defaultalg(A, b, assump::OperatorAssumptions{Bool})
                             eltype(b) <: Union{Float32, Float64, ComplexF32, ComplexF64}
                         DefaultAlgorithmChoice.AppleAccelerateLUFactorization
                     elseif (
-                            length(b) <= 100 || (isopenblas() && length(b) <= 500) ||
-                                (usemkl && length(b) <= 200)
+                            size(b, 1) <= 100 || (isopenblas() && size(b, 1) <= 500) ||
+                                (usemkl && size(b, 1) <= 200)
                         ) &&
                             (
                             A === nothing ? eltype(b) <: Union{Float32, Float64} :

src/factorization.jl

@@ -122,8 +122,10 @@ function _check_residual_safety(cache::LinearCache, alg, A_original, y)
     b = cache.b
     if cache.alg isa DefaultLinearSolver
         buf = cache.cacheval.residual_buf
-        if length(buf) != length(b)
-            resize!(buf, length(b))
+        if size(buf) != size(b)
+            # `resize!` only applies to vectors; matrix (batched) b just allocates.
+            buf = buf isa Vector && b isa AbstractVector ? resize!(buf, length(b)) :
+                similar(b)
         end
     else
         buf = similar(b)

src/iterative_wrappers.jl

@@ -276,6 +276,19 @@ function init_cacheval(
     return nothing
 end
 
+# Krylov workspaces only support vector right-hand sides. Batched (matrix) `b`
+# with a Krylov algorithm errors informatively at `init` time
+# (`_check_batched_rhs_support`); this method only exists so the default
+# polyalgorithm can still initialize its (unused) Krylov cacheval slots when a
+# factorization algorithm is chosen for a batched problem.
+function init_cacheval(
+        alg::LinearSolve.KrylovJL, A, b::AbstractMatrix, u, Pl, Pr,
+        maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
+        ::LinearSolve.OperatorAssumptions; zeroinit = true
+    )
+    return nothing
+end
+
 function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
     if cache.precsisfresh && !isnothing(alg.precs)
         Pl, Pr = alg.precs(cache.A, cache.p)

src/mkl.jl

@@ -386,7 +386,11 @@ function SciMLBase.solve!(
     if m > n
         Bc = copy(cache.b)
         getrs!('N', A.factors, A.ipiv, Bc; info)
-        copyto!(cache.u, 1, Bc, 1, n)
+        if cache.b isa AbstractMatrix
+            copyto!(cache.u, @view(Bc[1:n, :]))
+        else
+            copyto!(cache.u, 1, Bc, 1, n)
+        end
     else
         copyto!(cache.u, cache.b)
         getrs!('N', A.factors, A.ipiv, cache.u; info)
@@ -469,7 +473,11 @@ function SciMLBase.solve!(
     if m > n
         getrs!('N', A_lu.factors, A_lu.ipiv, b_32; info)
         # Convert back to original precision
-        cache.u[1:n] .= Torig.(b_32[1:n])
+        if cache.b isa AbstractMatrix
+            cache.u .= Torig.(@view(b_32[1:n, :]))
+        else
+            cache.u[1:n] .= Torig.(@view(b_32[1:n]))
+        end
     else
         copyto!(u_32, b_32)
         getrs!('N', A_lu.factors, A_lu.ipiv, u_32; info)