diff --git a/docs/src/release_notes.md b/docs/src/release_notes.md index 1a81d8d66..97902738a 100644 --- a/docs/src/release_notes.md +++ b/docs/src/release_notes.md @@ -8,6 +8,9 @@ - Dense `LUFactorization` refactorizations (`cache.A = X` then `solve!`) now reuse the cached pivot vector (and, without `alias_A`, the cached factors buffer) on Julia >= 1.11, making warm refactorization solves allocation-free. + - The dense `LUFactorization` pivot-buffer reuse with `alias_A = true` now also covers + the generic-kernel path (`NoPivot`/`RowNonZero` pivoting and non-BLAS element types), + on all supported Julia versions. ## v4.0 diff --git a/src/factorization.jl b/src/factorization.jl index 3f2a1052f..266c6c190 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -266,15 +266,18 @@ _get_residualsafety(alg::GenericLUFactorization) = alg.residualsafety # Dense-LU refactorization buffer reuse: `lu`/`lu!` allocate a fresh pivot # vector (and, without aliasing, a fresh factors copy) on every `cache.A = X` # refactorization. When the cached `LU`'s buffers match the incoming matrix, -# factorize through `LAPACK.getrf!` with the cached `ipiv` instead. The +# factorize through `LAPACK.getrf!` (or, outside the LAPACK fast path, the +# vendored `generic_lufact!`) with the cached `ipiv` instead. The # preallocated-`ipiv` `getrf!` method only exists on Julia >= 1.11; older -# releases keep the allocating path. +# releases keep the allocating LAPACK path, but `generic_lufact!` accepts a +# provided `ipiv` on every supported Julia version. +function _lu_cacheval_ipiv_matches(cacheval, A) + return cacheval isa LU && + cacheval.ipiv isa Vector{BlasInt} && + length(cacheval.ipiv) == min(size(A)...) +end @static if VERSION >= v"1.11" - function _reusable_lu_cacheval(cacheval, A) - return cacheval isa LU && - cacheval.ipiv isa Vector{BlasInt} && - length(cacheval.ipiv) == min(size(A)...) - end + _reusable_lu_cacheval(cacheval, A) = _lu_cacheval_ipiv_matches(cacheval, A) function _lu_reusing_ipiv!(A, ipiv::Vector{BlasInt}) factors, piv, info = LAPACK.getrf!(A, ipiv; check = false) return LU{eltype(factors), typeof(factors), typeof(piv)}( @@ -312,8 +315,19 @@ function SciMLBase.solve!(cache::LinearCache, alg::LUFactorization; kwargs...) # so refactorize in place and skip the O(n²) copy `lu` makes. pivot = _normalize_pivot(alg.pivot) if A isa StridedMatrix{<:LinearAlgebra.BlasFloat} && - pivot isa RowMaximum && _reusable_lu_cacheval(cacheval, A) - fact = _lu_reusing_ipiv!(A, cacheval.ipiv) + pivot isa RowMaximum + if _reusable_lu_cacheval(cacheval, A) + fact = _lu_reusing_ipiv!(A, cacheval.ipiv) + else + fact = lu!(A, pivot; check = false) + end + elseif A isa StridedMatrix && + pivot isa Union{RowMaximum, NoPivot, RowNonZero} && + _lu_cacheval_ipiv_matches(cacheval, A) + # `lu!` on a strided matrix outside the LAPACK fast path runs + # `generic_lufact!`, which would allocate a fresh pivot + # vector; call it directly with the cached one instead. + fact = generic_lufact!(A, pivot, cacheval.ipiv; check = false) else fact = lu!(A, pivot; check = false) end diff --git a/test/Core/lu_refactorization.jl b/test/Core/lu_refactorization.jl new file mode 100644 index 000000000..c74761c42 --- /dev/null +++ b/test/Core/lu_refactorization.jl @@ -0,0 +1,169 @@ +using LinearSolve, LinearAlgebra, StableRNGs, Test + +# Warm dense-LU refactorizations (`cache.A = X; solve!(cache)`) with +# `alias_A = true` must reuse the cacheval's pivot buffer instead of allocating +# a fresh `ipiv`/`LU` wrapper through `lu!` on every call. The LAPACK +# (`getrf!`) reuse path needs Julia >= 1.11; the generic-kernel path (NoPivot / +# non-BLAS eltypes) is allocation-free on every supported version. +const BLAS_IPIV_REUSE = VERSION >= v"1.11" + +rng = StableRNG(42) + +function refactor_solve!(cache, Awork, A) + copyto!(Awork, A) + cache.A = Awork + return solve!(cache) +end + +@testset "dense LU refactorization reuses pivot buffers" begin + n = 51 + A1 = rand(rng, n, n) + n * I + A2 = rand(rng, n, n) + n * I + b = rand(rng, n) + + # `maxalloc` is the per-refactorization allocation ceiling: 0 wherever the + # reuse path is active. On Julia 1.10 the BLAS (`getrf!`-with-`ipiv`) + # method does not exist, so those paths keep the allocating `lu!` (no + # assertion); the generic NoPivot kernel still reuses the pivot there, but + # 1.10's compiler does not always elide the small `LU`/solution wrapper + # constructions that 1.11+ removes, hence the small nonzero ceiling. + @testset "$(name)" for (name, alg, maxalloc) in ( + ("LUFactorization RowMaximum", LUFactorization(), BLAS_IPIV_REUSE ? 0 : nothing), + ( + "LUFactorization NoPivot", LUFactorization(pivot = NoPivot()), + VERSION >= v"1.11" ? 0 : 64, + ), + ("default algorithm", nothing, BLAS_IPIV_REUSE ? 0 : nothing), + ) + cache = init( + LinearProblem(copy(A1), copy(b)), alg; + alias = LinearAliasSpecifier(alias_A = true) + ) + Awork = cache.A + @test solve!(cache).u ≈ A1 \ b + + # correctness across repeated refactorize+solve cycles + for Ak in (A2, A1, A2) + sol = refactor_solve!(cache, Awork, Ak) + @test sol.retcode == ReturnCode.Success + @test sol.u ≈ Ak \ b + end + + # post-warmup refactorization must not allocate beyond the ceiling, + # and the cached pivot vector must be reused, not replaced + refactor_solve!(cache, Awork, A1) + ipiv_before = alg isa LUFactorization ? cache.cacheval.ipiv : nothing + alloc = @allocated refactor_solve!(cache, Awork, A2) + if maxalloc !== nothing + @test alloc <= maxalloc + if alg isa LUFactorization + @test cache.cacheval.ipiv === ipiv_before + end + end + @test cache.u ≈ A2 \ b + end +end + +@testset "singular refactorization reports Failure without throwing" begin + n = 8 + Agood = rand(rng, n, n) + n * I + Asing = copy(Agood) + Asing[:, 1] .= 0 # exactly rank-deficient, singular under any pivoting + b = rand(rng, n) + + @testset "$(name)" for (name, alg) in ( + ("RowMaximum", LUFactorization()), + ("NoPivot", LUFactorization(pivot = NoPivot())), + ) + cache = init( + LinearProblem(copy(Agood), copy(b)), alg; + alias = LinearAliasSpecifier(alias_A = true) + ) + Awork = cache.A + @test solve!(cache).retcode == ReturnCode.Success + + # warm cycle so the singular refactorization exercises the reuse path + refactor_solve!(cache, Awork, Agood) + @test refactor_solve!(cache, Awork, Asing).retcode == ReturnCode.Failure + + # the cache recovers on the next nonsingular refactorization + sol = refactor_solve!(cache, Awork, Agood) + @test sol.retcode == ReturnCode.Success + @test sol.u ≈ Agood \ b + end +end + +@testset "default algorithm QR safety fallback survives warm singular refactorization" begin + n = 8 + Agood = rand(rng, n, n) + n * I + Asing = copy(Agood) + Asing[:, 1] .= 0 + b = rand(rng, n) + + cache = init( + LinearProblem(copy(Agood), copy(b)), nothing; + alias = LinearAliasSpecifier(alias_A = true) + ) + Awork = cache.A + @test solve!(cache).retcode == ReturnCode.Success + refactor_solve!(cache, Awork, Agood) + + # LU fails on the singular A, so the default algorithm's safetyfallback + # rescues through column-pivoted QR (least-squares) instead of erroring + sol = refactor_solve!(cache, Awork, Asing) + @test sol.retcode == ReturnCode.Success + @test Asing * sol.u ≈ Asing * (qr(Asing, ColumnNorm()) \ b) + + sol = refactor_solve!(cache, Awork, Agood) + @test sol.retcode == ReturnCode.Success + @test sol.u ≈ Agood \ b +end + +@testset "generic (non-BLAS) eltype reuses the cached pivot" begin + n = 6 + A1 = big.(rand(rng, n, n)) + n * I + A2 = big.(rand(rng, n, n)) + n * I + b = big.(rand(rng, n)) + + cache = init( + LinearProblem(copy(A1), copy(b)), LUFactorization(); + alias = LinearAliasSpecifier(alias_A = true) + ) + Awork = cache.A + @test solve!(cache).u ≈ A1 \ b + for Ak in (A2, A1, A2) + sol = refactor_solve!(cache, Awork, Ak) + @test sol.retcode == ReturnCode.Success + @test sol.u ≈ Ak \ b + end +end + +@testset "size change between refactorizations reallocates the pivot" begin + n1, n2 = 5, 9 + A1 = rand(rng, n1, n1) + n1 * I + b1 = rand(rng, n1) + + cache = init( + LinearProblem(copy(A1), copy(b1)), LUFactorization(); + alias = LinearAliasSpecifier(alias_A = true) + ) + @test solve!(cache).u ≈ A1 \ b1 + + resize!(cache, n2) + A2 = rand(rng, n2, n2) + n2 * I + b2 = rand(rng, n2) + cache.A = copy(A2) + cache.b = copy(b2) + cache.u = zeros(n2) + @test solve!(cache).u ≈ A2 \ b2 + + # warm cycles at the new size are allocation-free again + Awork = cache.A + refactor_solve!(cache, Awork, A2) + refactor_solve!(cache, Awork, A2) + alloc = @allocated refactor_solve!(cache, Awork, A2) + if BLAS_IPIV_REUSE + @test alloc == 0 + end + @test cache.u ≈ A2 \ b2 +end diff --git a/test/runtests.jl b/test/runtests.jl index c44370b6a..f49ff6797 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -64,6 +64,7 @@ else @time @safetestset "Basic Tests" include("Core/basictests.jl") @time @safetestset "Batched RHS" include("Core/batch.jl") @time @safetestset "GESV Factorization" include("Core/gesv.jl") + @time @safetestset "LU Refactorization Reuse" include("Core/lu_refactorization.jl") @time @safetestset "Return codes" include("Core/retcodes.jl") @time @safetestset "Re-solve" include("Core/resolve.jl") @time @safetestset "Zero Initialization Tests" include("Core/zeroinittests.jl")