Preallocate GenericLU packing workspace

src/SupernodalLU/numeric.jl

@@ -315,6 +315,9 @@ end
 # or above `dense_threshold` is a `MethodError`.  `hasfield` on a concrete type
 # folds at compile time, so the check costs nothing at runtime.
 function _lu_from_cacheval(cv)
+    if cv isa LinearSolve._GenericLUFactorizationCache
+        return cv.fact
+    end
     if hasfield(typeof(cv), :factors) && hasfield(typeof(cv), :ipiv)
         return cv
     end

src/adjoint_factorization.jl

@@ -98,6 +98,8 @@ end
 function _standard_cache_factorization(cacheval)
     if cacheval isa Factorization
         return cacheval
+    elseif cacheval isa _GenericLUFactorizationCache
+        return cacheval.fact
     elseif cacheval isa Tuple && !isempty(cacheval) && first(cacheval) isa Factorization
         return first(cacheval)
     else

src/blocked_lufact.jl

@@ -14,6 +14,22 @@ const _BLOCKED_LU_ROWBLOCK = 384
 # amortize panel/trsm overhead once the Schur update dominates.
 _blocked_lu_default_panel(minmn::Int) = minmn <= 160 ? 8 : 16
 
+function _blocked_lu_pack_size(m::Int, n::Int, minmn::Int, nb::Int)
+    len = 0
+    j0 = 1
+    while j0 <= minmn
+        jb = min(nb, minmn - j0 + 1)
+        j1 = j0 + jb - 1
+        rows = m - j1
+        if rows >= 64 && n - j1 >= 32
+            ldp = rows % 256 == 0 ? rows + 4 : rows
+            len = max(len, ldp * jb)
+        end
+        j0 += jb
+    end
+    return len
+end
+
 # Row-maximum pivot search in two passes: a `>`-select max reduction that
 # vectorizes (NaN compares false, so NaNs are ignored exactly like the
 # scalar stdlib/LAPACK search), then first-index-of-max, matching the
@@ -305,7 +321,6 @@ function _blocked_lu_schur_packed!(
     rows = m - i0 + 1
     jb = j1 - j0 + 1
     ldp = rows % 256 == 0 ? rows + 4 : rows
-    length(pack) < ldp * jb && resize!(pack, ldp * jb)
     _blocked_lu_pack_panel!(pack, A, i0, m, j0, j1, ldp)
     @inbounds begin
         ib = i0
@@ -582,7 +597,6 @@ function _blocked_lu_schur_micro!(
     rows = m - i0 + 1
     jb = j1 - j0 + 1
     ldp = rows % 256 == 0 ? rows + 4 : rows
-    length(pack) < ldp * jb && resize!(pack, ldp * jb)
     _blocked_lu_pack_panel!(pack, A, i0, m, j0, j1, ldp)
     V = _blocked_lu_vectype(T)
     # Must equal the tile's row step; a mismatch silently double-applies or
@@ -643,9 +657,8 @@ end
 
 function _blocked_lufact!(
         A::AbstractMatrix{T}, ipiv, m::Int, n::Int, minmn::Int,
-        nb::Int, rowblock::Int
+        nb::Int, rowblock::Int, pack::Vector{T}
     ) where {T}
-    pack = T[]
     info = 0
     j0 = 1
     while j0 <= minmn
@@ -665,6 +678,14 @@ function _blocked_lufact!(
     return info
 end
 
+function _blocked_lufact!(
+        A::AbstractMatrix{T}, ipiv, m::Int, n::Int, minmn::Int,
+        nb::Int, rowblock::Int
+    ) where {T}
+    pack = Vector{T}(undef, _blocked_lu_pack_size(m, n, minmn, nb))
+    return _blocked_lufact!(A, ipiv, m, n, minmn, nb, rowblock, pack)
+end
+
 # The `GenericLUFactorization` fast path: real strided float matrices with
 # `RowMaximum` pivoting take the blocked kernel; everything else falls through
 # to the scalar method. Semantics match `generic_lufact!` with a provided
@@ -676,6 +697,26 @@ function generic_lufact!(
         ipiv::AbstractVector{<:Integer};
         check::Bool = true, allowsingular::Bool = false
     ) where {T <: Union{Float32, Float64}}
+    return _blocked_generic_lufact!(
+        A, pivot, ipiv, nothing; check = check, allowsingular = allowsingular
+    )
+end
+
+function generic_lufact!(
+        A::StridedMatrix{T}, pivot::RowMaximum,
+        ipiv::AbstractVector{<:Integer}, pack::Vector{T};
+        check::Bool = true, allowsingular::Bool = false
+    ) where {T <: Union{Float32, Float64}}
+    return _blocked_generic_lufact!(
+        A, pivot, ipiv, pack; check = check, allowsingular = allowsingular
+    )
+end
+
+function _blocked_generic_lufact!(
+        A::StridedMatrix{T}, pivot::RowMaximum,
+        ipiv::AbstractVector{<:Integer}, pack::Union{Nothing, Vector{T}};
+        check::Bool = true, allowsingular::Bool = false
+    ) where {T <: Union{Float32, Float64}}
     Base.require_one_based_indexing(A, ipiv)
     if check && !all(isfinite, A)
         throw(ArgumentError("matrix contains Infs or NaNs"))
@@ -687,9 +728,11 @@ function generic_lufact!(
     info = if minmn <= _BLOCKED_LU_UNBLOCKED_CUTOFF
         _blocked_lu_unblocked!(A, ipiv, m, n)
     else
+        nb = _blocked_lu_default_panel(minmn)
+        pack === nothing &&
+            (pack = Vector{T}(undef, _blocked_lu_pack_size(m, n, minmn, nb)))
         _blocked_lufact!(
-            A, ipiv, m, n, minmn, _blocked_lu_default_panel(minmn),
-            _BLOCKED_LU_ROWBLOCK
+            A, ipiv, m, n, minmn, nb, _BLOCKED_LU_ROWBLOCK, pack
         )
     end
     check && !allowsingular && info > 0 &&

src/common.jl

@@ -294,6 +294,7 @@ end
                 cache.cacheval.a_backup_synced = false
             end
         end
+        update_cacheval!(cache, :A, x)
     elseif name === :p
         setfield!(cache, :precsisfresh, true)
     elseif name === :b

src/default.jl

@@ -47,6 +47,7 @@ mutable struct DefaultLinearSolverInit{
 end
 
 function resize_cacheval!(cache, cacheval::DefaultLinearSolverInit, i)
+    resize_cacheval!(cache, cacheval.GenericLUFactorization, i)
     A_backup = cacheval.A_backup
     return if A_backup isa AbstractMatrix
         setfield!(cacheval, :A_backup, similar(A_backup, i, i))
@@ -55,6 +56,11 @@ function resize_cacheval!(cache, cacheval::DefaultLinearSolverInit, i)
     end
 end
 
+function update_cacheval!(cache, cacheval::DefaultLinearSolverInit, name::Symbol, A)
+    name === :A && update_cacheval!(cache, cacheval.GenericLUFactorization, name, A)
+    return cacheval
+end
+
 @generated function __setfield!(cache::DefaultLinearSolverInit, alg::DefaultLinearSolver, v)
     ex = :()
     for alg in first.(EnumX.symbol_map(DefaultAlgorithmChoice.T))
@@ -1192,15 +1198,14 @@ end
 @generated function defaultalg_adjoint_eval(cache::LinearCache, dy)
     ex = :()
     for alg in first.(EnumX.symbol_map(DefaultAlgorithmChoice.T))
-        newex = if alg in Symbol.(
-                (
-                    DefaultAlgorithmChoice.RFLUFactorization,
-                    DefaultAlgorithmChoice.GenericLUFactorization,
-                )
-            )
+        newex = if alg == Symbol(DefaultAlgorithmChoice.RFLUFactorization)
             quote
                 getproperty(cache.cacheval, $(Meta.quot(alg)))[1]' \ dy
             end
+        elseif alg == Symbol(DefaultAlgorithmChoice.GenericLUFactorization)
+            quote
+                getproperty(cache.cacheval, $(Meta.quot(alg))).fact' \ dy
+            end
         elseif alg == Symbol(DefaultAlgorithmChoice.MKLLUFactorization)
             quote
                 A = getproperty(cache.cacheval, $(Meta.quot(alg)))[1]

src/factorization.jl

@@ -299,6 +299,69 @@ end
 
 GenericLUFactorization(pivot = RowMaximum(); residualsafety::Bool = false) = GenericLUFactorization(pivot, residualsafety)
 
+mutable struct _GenericLUFactorizationCache{F, I, W}
+    fact::F
+    ipiv::I
+    workspace::W
+end
+
+_generic_lu_workspace(A, pivot) = nothing
+
+function _generic_lu_workspace(A::StridedMatrix{T}, ::RowMaximum) where {T <: Union{Float32, Float64}}
+    m, n = size(A)
+    minmn = min(m, n)
+    nb = _blocked_lu_default_panel(minmn)
+    return Vector{T}(undef, _blocked_lu_pack_size(m, n, minmn, nb))
+end
+
+function _generic_lufact!(A, pivot, ipiv, ::Nothing; kwargs...)
+    return generic_lufact!(A, pivot, ipiv; kwargs...)
+end
+
+function _generic_lufact!(A, pivot, ipiv, workspace::Vector; kwargs...)
+    return generic_lufact!(A, pivot, ipiv, workspace; kwargs...)
+end
+
+function _generic_lu_solve!(cacheval, A, u, b, pivot, isfresh::Bool)
+    if isfresh
+        fact = _generic_lufact!(
+            A, pivot, cacheval.ipiv, cacheval.workspace; check = false
+        )
+        cacheval.fact = fact
+        LinearAlgebra.issuccess(fact) || return false
+    end
+    fact = cacheval.fact
+    if fact isa LinearAlgebra.LU
+        _generic_lu_ldiv!(u, fact, b)
+    else
+        ldiv!(u, fact, b)
+    end
+    return true
+end
+
+_resize_generic_lu_workspace!(::Nothing, m::Int, n::Int) = nothing
+
+function _resize_generic_lu_workspace!(workspace::Vector, m::Int, n::Int)
+    minmn = min(m, n)
+    nb = _blocked_lu_default_panel(minmn)
+    resize!(workspace, _blocked_lu_pack_size(m, n, minmn, nb))
+    return nothing
+end
+
+function _resize_generic_lu_cache!(cacheval::_GenericLUFactorizationCache, m::Int, n::Int)
+    resize!(cacheval.ipiv, min(m, n))
+    _resize_generic_lu_workspace!(cacheval.workspace, m, n)
+    return nothing
+end
+
+resize_cacheval!(cache, cacheval::_GenericLUFactorizationCache, i) =
+    _resize_generic_lu_cache!(cacheval, i, i)
+
+function update_cacheval!(cache, cacheval::_GenericLUFactorizationCache, name::Symbol, A)
+    name === :A && _resize_generic_lu_cache!(cacheval, size(A, 1), size(A, 2))
+    return cacheval
+end
+
 # Pure-Julia LU back-solve used by GenericLUFactorization. A pivoted LU vector
 # solve at small N is a handful of flops but ~290 ns through OpenBLAS/MKL
 # getrs! (call overhead); the same algorithm written out as scalar loops is
@@ -630,26 +693,29 @@ function init_cacheval(
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
+    A = convert(AbstractMatrix, A)
     ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...))
-    # `solve!` stores `(generic_lufact!(A, ...), ipiv)`, a `LinearAlgebra.LU` whose
-    # `factors` is `convert(AbstractMatrix, A)` and whose pivot is this
-    # `Vector{BlasInt}`. `lu_instance` would type the pivot after `A`'s container
-    # (e.g. a `FixedSizeVector` for a `FixedSizeArray`), so rebuild the instance
-    # with the `Vector` pivot to keep the cacheval slot type matching.
-    luinst = ArrayInterface.lu_instance(convert(AbstractMatrix, A))
+    # `lu_instance` may type its pivot after `A`'s container, so rebuild stdlib
+    # `LU` instances with the cache-owned `Vector{BlasInt}` pivot.
+    luinst = ArrayInterface.lu_instance(A)
     # `lu_instance` may return a non-`LinearAlgebra.LU` (e.g. `StaticArrays.LU`
     # for a `SizedMatrix`, with fields `L`/`U`/`p` rather than `factors`/`info`);
     # those already carry a `Vector` pivot, so use them as-is.
-    luinst isa LinearAlgebra.LU || return luinst, ipiv
-    return LinearAlgebra.LU(luinst.factors, ipiv, luinst.info), ipiv
+    workspace = _generic_lu_workspace(A, alg.pivot)
+    luinst isa LinearAlgebra.LU ||
+        return _GenericLUFactorizationCache(luinst, ipiv, workspace)
+    fact = LinearAlgebra.LU(luinst.factors, ipiv, luinst.info)
+    return _GenericLUFactorizationCache(fact, ipiv, workspace)
 end
 
 function init_cacheval(
         alg::GenericLUFactorization, A::Matrix{Float64}, b, u, Pl, Pr,
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
-    return PREALLOCATED_LU, PREALLOCATED_IPIV
+    ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...))
+    workspace = _generic_lu_workspace(A, alg.pivot)
+    return _GenericLUFactorizationCache(PREALLOCATED_LU, ipiv, workspace)
 end
 
 function SciMLBase.solve!(
@@ -662,31 +728,15 @@ function SciMLBase.solve!(
     needs_backup = check_safety ||
         (cache.alg isa DefaultLinearSolver && cache.alg.safetyfallback && cache.isfresh)
     A_original = needs_backup ? _copy_A_for_safety(cache) : A
-    fact, ipiv = LinearSolve.@get_cacheval(cache, :GenericLUFactorization)
-
-    if cache.isfresh
-        if length(ipiv) != min(size(A)...)
-            ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...))
-        end
-        fact = generic_lufact!(A, alg.pivot, ipiv; check = false)
-        cache.cacheval = (fact, ipiv)
+    cacheval = LinearSolve.@get_cacheval(cache, :GenericLUFactorization)
 
-        if !LinearAlgebra.issuccess(fact)
-            return SciMLBase.build_linear_solution(
-                alg, cache.u, nothing, nothing; retcode = ReturnCode.Failure
-            )
-        end
-
-        cache.isfresh = false
-    end
-    F = LinearSolve.@get_cacheval(cache, :GenericLUFactorization)[1]
-    # Prefer the pure-Julia back-solve for `LinearAlgebra.LU` (the common
-    # cacheval). Non-stdlib LU types (e.g. StaticArrays) keep their own `ldiv!`.
-    y = if F isa LinearAlgebra.LU
-        _generic_lu_ldiv!(cache.u, F, cache.b)
-    else
-        ldiv!(cache.u, F, cache.b)
+    if !_generic_lu_solve!(cacheval, A, cache.u, cache.b, alg.pivot, cache.isfresh)
+        return SciMLBase.build_linear_solution(
+            alg, cache.u, nothing, nothing; retcode = ReturnCode.Failure
+        )
     end
+    cache.isfresh = false
+    y = cache.u
 
     if check_safety
         failed = _check_residual_safety(cache, alg, A_original, y)
@@ -720,8 +770,9 @@ function init_cacheval(
     )
     error_no_cudss_lu(A)
     A isa GPUArraysCore.AnyGPUArray && return nothing
-    ipiv = Vector{LinearAlgebra.BlasInt}(undef, 0)
-    return LinearAlgebra.generic_lufact!(_typed_copy(A), alg.pivot; check = false), ipiv
+    ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...))
+    fact = LinearAlgebra.generic_lufact!(_typed_copy(A), alg.pivot; check = false)
+    return _GenericLUFactorizationCache(fact, ipiv, nothing)
 end
 
 const PREALLOCATED_LU = ArrayInterface.lu_instance(rand(1, 1))

test/Core/genericlu_naive_ldiv.jl

@@ -90,7 +90,7 @@ end
         b = rand(n)
         cache = init(LinearProblem(mk(copy(A)), copy(b)), GenericLUFactorization())
         sol = solve!(cache)
-        F = first(cache.cacheval)
+        F = cache.cacheval.fact
         xldiv = copy(b)
         ldiv!(F, xldiv)
         @test sol.u != xldiv
@@ -162,7 +162,7 @@ end
                 @test SciMLBase.successful_retcode(sol)
                 @test sol.u ≈ xref rtol = rtol * n
 
-                F = first(cache.cacheval)
+                F = cache.cacheval.fact
                 @test F.factors isa (wrap === adjoint ? Adjoint : Transpose)
 
                 x1 = copy(b)

test/qa/allocations.jl

@@ -16,6 +16,43 @@ end
     return info
 end
 
+@check_allocs function allocation_checked_generic_lu_solve!(cache)
+    success = LinearSolve._generic_lu_solve!(
+        cache.cacheval, cache.A, cache.u, cache.b, cache.alg.pivot, cache.isfresh
+    )
+    cache.isfresh = !success
+    return success
+end
+
+function generic_lu_solve_allocations(cache, Awork, A)
+    copyto!(Awork, A)
+    cache.A = Awork
+    solve!(cache)
+    copyto!(Awork, A)
+    cache.A = Awork
+    return @allocated solve!(cache)
+end
+
+@testset "GenericLUFactorization solve! is allocation-free" begin
+    n = 100
+    A = rand(n, n) + n * I
+    b = rand(n)
+    cache = init(LinearProblem(copy(A), b), GenericLUFactorization())
+
+    A2 = rand(n, n) + n * I
+    if VERSION >= v"1.12"
+        @test generic_lu_solve_allocations(cache, cache.A, A2) == 0
+    else
+        generic_lu_solve_allocations(cache, cache.A, A2)
+    end
+    @test cache.u ≈ A2 \ b
+
+    copyto!(cache.A, A2)
+    cache.A = cache.A
+    @test allocation_checked_generic_lu_solve!(cache)
+    @test cache.u ≈ A2 \ b
+end
+
 function test_allocation_free_refactorization(alg, ::Type{T}) where {T}
     A1 = T[4 1; 2 3]
     A2 = T[3 -1; 1 2]