From a74d1fd90a20b6b216ef74c3b5a3391273483027 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 4 Jul 2026 16:53:16 -0400 Subject: [PATCH 1/2] Batched RHS via Krylov.jl block methods for KrylovJL_GMRES/MINRES Krylov.jl provides block_gmres and block_minres, so those KrylovJL variants now accept matrix right-hand sides natively via BlockGmresWorkspace/BlockMinresWorkspace instead of erroring; the remaining Krylov methods (no block variants) keep the informative ArgumentError. The DefaultLinearSolver operator path's GMRES choice is un-gated accordingly, so `solve(LinearProblem(op, B))` with an operator A and matrix B now works out of the box. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Fable 5 --- src/common.jl | 20 +++++++++------- src/iterative_wrappers.jl | 48 +++++++++++++++++++++++++++++++++------ 2 files changed, 53 insertions(+), 15 deletions(-) diff --git a/src/common.jl b/src/common.jl index fe5084147..00d132dd3 100644 --- a/src/common.jl +++ b/src/common.jl @@ -448,21 +448,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 diff --git a/src/iterative_wrappers.jl b/src/iterative_wrappers.jl index dd203ca22..efda88eb9 100644 --- a/src/iterative_wrappers.jl +++ b/src/iterative_wrappers.jl @@ -276,9 +276,10 @@ 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( @@ -286,9 +287,30 @@ function init_cacheval( 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( From 620b391fd5ddbab675fcd41a8ffb48c101a55501 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 5 Jul 2026 02:21:51 -0400 Subject: [PATCH 2/2] alias_A=true enables in-place dense LU refactorization (v4.1.0) The resolved LinearAliasSpecifier.alias_A from init is now stored on the LinearCache. When the user permitted overwriting A, the dense LUFactorization refactorization path (after cache.A = X) runs lu!(cache.A) in place instead of paying lu's O(n^2) copy on every refactorization. Also updates the stale batch testset that still expected KrylovJL_GMRES to reject matrix right-hand sides: block GMRES/MINRES support made that path solve instead of throw, so it now asserts correct batched solutions and keeps the informative-error checks on the methods without block variants. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Fable 5 --- Project.toml | 2 +- docs/src/basics/common_solver_opts.md | 7 +++ ext/LinearSolveHYPREExt.jl | 2 +- src/common.jl | 10 ++++- src/factorization.jl | 6 +++ test/Core/batch.jl | 61 +++++++++++++++++++++++++-- 6 files changed, 81 insertions(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index 7278a431a..f7f4d7cc2 100644 --- a/Project.toml +++ b/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] diff --git a/docs/src/basics/common_solver_opts.md b/docs/src/basics/common_solver_opts.md index ce03c2701..80eca1552 100644 --- a/docs/src/basics/common_solver_opts.md +++ b/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). diff --git a/ext/LinearSolveHYPREExt.jl b/ext/LinearSolveHYPREExt.jl index 94f3e47d2..b8ec68ada 100644 --- a/ext/LinearSolveHYPREExt.jl +++ b/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 diff --git a/src/common.jl b/src/common.jl index 00d132dd3..bb24eebe4 100644 --- a/src/common.jl +++ b/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) @@ -617,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 diff --git a/src/factorization.jl b/src/factorization.jl index 488ac0960..a7df7008c 100644 --- a/src/factorization.jl +++ b/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 diff --git a/test/Core/batch.jl b/test/Core/batch.jl index ab25ad7b9..0abc4cb85 100644 --- a/test/Core/batch.jl +++ b/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