Complete allocation-free dense LU refactorization: generic-kernel pivot reuse + regression tests

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
 

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

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

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")