Batched RHS via Krylov.jl block methods (GMRES/MINRES)

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearSolve"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
-version = "4.0.0"
+version = "4.1.0"
 authors = ["SciML"]
 
 [deps]

docs/src/basics/common_solver_opts.md

@@ -12,6 +12,13 @@ The following are the options these algorithms take, along with their defaults.
     `A` and `b` can be written to and changed by the solver algorithm. When fields are `nothing`
     the default behavior is used, which is to default to `true` when the algorithm is known
     not to modify the matrices, and false otherwise.
+
+    `alias_A = true` also extends to caching interface refactorizations: after replacing
+    the matrix via `cache.A = A2`, dense factorizations such as `LUFactorization` will
+    factorize in place (`lu!(cache.A)`), overwriting `cache.A` with its factors instead
+    of making an O(n²) copy on every refactorization. Only opt in when you own the matrix
+    handed to the cache and refill it before each new solve, since its contents are
+    destroyed by the factorization.
   - `verbose`: Whether to print extra information. Defaults to `false`.
   - `assumptions`: Sets the assumptions of the operator in order to effect the default
     choice algorithm. See the [Operator Assumptions page for more details](@ref assumptions).

ext/LinearSolveHYPREExt.jl

@@ -201,7 +201,7 @@ function SciMLBase.init(
         typeof(__issquare(assumptions)), typeof(sensealg), Nothing,
     }(
         A, b, u0, p, alg, cacheval, isfresh, precsisfresh, Pl, Pr, abstol, reltol,
-        maxiters, verb_spec, assumptions, sensealg, nothing
+        maxiters, verb_spec, assumptions, sensealg, nothing, alias_A
     )
     return cache
 end

src/common.jl

@@ -210,6 +210,10 @@ solving and caching of factorizations and intermediate results.
 - `verbose::LinearVerbosity`: Whether to print verbose output during solving.
 - `assumptions::OperatorAssumptions{issq}`: Assumptions about the operator properties.
 - `sensealg::S`: Sensitivity analysis algorithm for automatic differentiation.
+- `alias_A::Bool`: The resolved `LinearAliasSpecifier.alias_A` from `init`. When `true`,
+  the user has permitted LinearSolve to overwrite `A`; dense factorizations may then
+  refactorize in place (e.g. `lu!(cache.A)`) after `cache.A` is replaced, skipping the
+  O(n²) copy.
 
 ## Usage
 
@@ -244,6 +248,10 @@ mutable struct LinearCache{TA, Tb, Tu, Tp, Talg, Tc, Tl, Tr, Ttol, Tlv <: Linear
     # factorizations (`nothing` otherwise; the default solver carries its own in
     # `DefaultLinearSolverInit`). Set once at `init`; persists across `reinit!`.
     sparse_reduction::Tred
+    # Resolved `LinearAliasSpecifier.alias_A` from `init` (defaults applied, so
+    # never `nothing`). `true` means the user permitted overwriting `A`, which
+    # also permits in-place refactorization (e.g. `lu!(A)`) after `cache.A = X`.
+    alias_A::Bool
 end
 
 function Base.setproperty!(cache::LinearCache, name::Symbol, x)
@@ -448,21 +456,25 @@ function _check_batched_rhs_support(alg::AbstractKrylovSubspaceMethod, b::Abstra
     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()`)."
+                "algorithms and block Krylov methods; $(nameof(typeof(alg))) supports " *
+                "only vector `b`. Use KrylovJL_GMRES/KrylovJL_MINRES (block methods), " *
+                "a factorization algorithm (e.g. `LUFactorization()`), or solve " *
+                "column-by-column."
         )
     )
 end
 function _check_batched_rhs_support(alg::DefaultLinearSolver, b::AbstractMatrix)
-    if alg.alg === DefaultAlgorithmChoice.KrylovJL_GMRES ||
-            alg.alg === DefaultAlgorithmChoice.KrylovJL_CRAIGMR ||
+    # KrylovJL_GMRES is fine: it dispatches to Krylov.jl's block GMRES for
+    # matrix b. CRAIGMR/LSMR (least-squares operator defaults) have no block
+    # variants.
+    if 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."
+                "Batched (matrix) right-hand sides are not supported by the " *
+                    "least-squares Krylov method $(alg.alg) the default algorithm " *
+                    "selected for this operator. Solve column-by-column or use a " *
+                    "factorization algorithm."
             )
         )
     end
@@ -613,7 +625,7 @@ function __init(
         typeof(sensealg), typeof(sparse_reduction),
     }(
         A, b, u0_, p, alg, cacheval, isfresh, precsisfresh, Pl, Pr, abstol, reltol,
-        maxiters, verbose_spec, assumptions, sensealg, sparse_reduction
+        maxiters, verbose_spec, assumptions, sensealg, sparse_reduction, alias_A
     )
     return cache
 end

src/factorization.jl

@@ -282,6 +282,12 @@ function SciMLBase.solve!(cache::LinearCache, alg::LUFactorization; kwargs...)
                 else
                     fact = lu!(cacheval, A, check = false)
                 end
+            elseif cache.alias_A && !issparsematrix(A) &&
+                    !(A isa GPUArraysCore.AnyGPUArray) &&
+                    ArrayInterface.can_setindex(typeof(A))
+                # The user permitted overwriting A (`alias_A = true` at `init`),
+                # so refactorize in place and skip the O(n²) copy `lu` makes.
+                fact = lu!(A, _normalize_pivot(alg.pivot); check = false)
             else
                 fact = lu(A, check = false)
             end

src/iterative_wrappers.jl

@@ -276,19 +276,41 @@ 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
+# Batched (matrix) right-hand sides: Krylov.jl provides block methods for GMRES
+# and MINRES, so those get real block workspaces; the other methods have no
+# block variant and error informatively at `init` time
+# (`_check_batched_rhs_support`). The `nothing` fallback 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
     )
+    if alg.KrylovAlg === Krylov.gmres!
+        return Krylov.BlockGmresWorkspace(A, b)
+    elseif alg.KrylovAlg === Krylov.minres!
+        return Krylov.BlockMinresWorkspace(A, b)
+    end
     return nothing
 end
 
+# Krylov.jl provides block methods for GMRES and MINRES, so those KrylovJL
+# variants support batched right-hand sides natively (via BlockGmresWorkspace /
+# BlockMinresWorkspace); the remaining Krylov methods have no block variant.
+function _check_batched_rhs_support(alg::KrylovJL, b::AbstractMatrix)
+    (alg.KrylovAlg === Krylov.gmres! || alg.KrylovAlg === Krylov.minres!) &&
+        return nothing
+    throw(
+        ArgumentError(
+            "$(nameof(typeof(alg))) with $(alg.KrylovAlg) supports only vector `b`: " *
+                "Krylov.jl provides block (batched) methods only for GMRES and MINRES. " *
+                "Use KrylovJL_GMRES/KrylovJL_MINRES, a factorization algorithm, or " *
+                "solve column-by-column."
+        )
+    )
+end
+
 function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
     if cache.precsisfresh && !isnothing(alg.precs)
         Pl, Pr = alg.precs(cache.A, cache.p)
@@ -363,6 +385,15 @@ function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
             verbose, :no_right_preconditioning
         )
         Krylov.krylov_solve!(args...; M, kwargs...)
+    elseif cache.cacheval isa Krylov.BlockGmresWorkspace
+        Krylov.krylov_solve!(args...; M, N, restart = alg.gmres_restart > 0, kwargs...)
+    elseif cache.cacheval isa Krylov.BlockMinresWorkspace
+        N !== I &&
+            @SciMLMessage(
+            "$(alg.KrylovAlg) doesn't support right preconditioning.",
+            verbose, :no_right_preconditioning
+        )
+        Krylov.krylov_solve!(args...; M, kwargs...)
     elseif cache.cacheval isa Krylov.LsmrWorkspace ||
             cache.cacheval isa Krylov.LsqrWorkspace ||
             cache.cacheval isa Krylov.LslqWorkspace
@@ -398,12 +429,15 @@ function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
         ReturnCode.Success
     end
 
-    # Copy the solution to the allocated output vector
+    # Copy the solution to the allocated output array (block workspaces store
+    # the batched solution in `X` rather than `x`)
     cacheval = @get_cacheval(cache, :KrylovJL_GMRES)
-    if cache.u !== cacheval.x && ArrayInterface.can_setindex(cache.u)
-        cache.u .= cacheval.x
+    xsol = cacheval isa Union{Krylov.BlockGmresWorkspace, Krylov.BlockMinresWorkspace} ?
+        cacheval.X : cacheval.x
+    if cache.u !== xsol && ArrayInterface.can_setindex(cache.u)
+        cache.u .= xsol
     else
-        cache.u = convert(typeof(cache.u), cacheval.x)
+        cache.u = convert(typeof(cache.u), xsol)
     end
 
     return SciMLBase.build_linear_solution(

test/Core/batch.jl

@@ -146,19 +146,31 @@ end
     @test sol.u isa Matrix
 end
 
-@testset "Iterative algorithms error informatively" begin
+@testset "Block Krylov methods solve batched RHS" begin
+    A = rand(n, n) + n * I
+    B = rand(n, k)
+    sol = solve(LinearProblem(A, B), KrylovJL_GMRES())
+    @test SciMLBase.successful_retcode(sol)
+    @test sol.u ≈ A \ B
+
+    S = A + A' + 2n * I # symmetric for MINRES
+    sol = solve(LinearProblem(S, B), KrylovJL_MINRES())
+    @test SciMLBase.successful_retcode(sol)
+    @test sol.u ≈ S \ B
+end
+
+@testset "Iterative algorithms without block variants error informatively" begin
     A = rand(n, n) + n * I
     B = rand(n, k)
-    @test_throws ArgumentError solve(LinearProblem(A, B), KrylovJL_GMRES())
     @test_throws ArgumentError solve(LinearProblem(A, B), SimpleGMRES())
     @test_throws ArgumentError solve(LinearProblem(A, B), KrylovJL_CG())
     @test_throws ArgumentError solve(LinearProblem(A, B), SimpleLUFactorization())
     err = try
-        solve(LinearProblem(A, B), KrylovJL_GMRES())
+        solve(LinearProblem(A, B), KrylovJL_CG())
     catch e
         e
     end
-    @test occursin("Batched", err.msg)
+    @test occursin("batched", err.msg)
     @test occursin("KrylovJL", err.msg)
 end
 
@@ -178,3 +190,44 @@ end
     @test size(solmat.u) == (n, 1)
     @test vec(solmat.u) ≈ solvec.u
 end
+
+@testset "alias_A = true: LUFactorization refactorizes in place" begin
+    na = 100
+    D1 = rand(na, na) + na * I
+    D2 = rand(na, na) + na * I
+    B = rand(na, k)
+
+    function realias_solve!(cache, Abuf, D)
+        copyto!(Abuf, D)
+        cache.A = Abuf
+        return solve!(cache)
+    end
+
+    # alias_A = true: the user permitted overwriting A, so refactorization runs
+    # lu! on cache.A directly; repeated solves stay correct as long as the
+    # caller refills A between solves.
+    Abuf = copy(D1)
+    cache = init(
+        LinearProblem(Abuf, copy(B)), LUFactorization();
+        alias = LinearAliasSpecifier(alias_A = true)
+    )
+    @test cache.A === Abuf # not copied at init
+    @test solve!(cache).u ≈ D1 \ B
+    @test !(Abuf ≈ D1) # was overwritten by its LU factors
+    sol = realias_solve!(cache, Abuf, D2)
+    @test sol.u ≈ D2 \ B
+    @test cache.A === Abuf
+
+    # Reuse must not allocate a fresh n×n copy: O(n) small allocations only.
+    realias_solve!(cache, Abuf, D1) # warm up this exact path
+    alloc = @allocated realias_solve!(cache, Abuf, D2)
+    @test alloc < 20000 # n^2 * 8 = 80000 bytes would mean A was copied
+
+    # default (alias_A = false): cache.A is left untouched by refactorization.
+    cache = init(LinearProblem(copy(D1), copy(B)), LUFactorization())
+    @test solve!(cache).u ≈ D1 \ B
+    @test cache.A ≈ D1
+    cache.A = copy(D2)
+    @test solve!(cache).u ≈ D2 \ B
+    @test cache.A ≈ D2
+end