From 3ceacdcf15451c4c212358f053a59dd11c6adb97 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 10 Aug 2026 02:12:34 -0400 Subject: [PATCH] Preallocate GenericLU packing workspace Co-Authored-By: Chris Rackauckas --- src/SupernodalLU/numeric.jl | 3 + src/adjoint_factorization.jl | 2 + src/blocked_lufact.jl | 55 ++++++++++++-- src/common.jl | 1 + src/default.jl | 17 +++-- src/factorization.jl | 119 +++++++++++++++++++++--------- test/Core/genericlu_naive_ldiv.jl | 4 +- test/qa/allocations.jl | 37 ++++++++++ 8 files changed, 190 insertions(+), 48 deletions(-) diff --git a/src/SupernodalLU/numeric.jl b/src/SupernodalLU/numeric.jl index 565f21362..9c32a452c 100644 --- a/src/SupernodalLU/numeric.jl +++ b/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 diff --git a/src/adjoint_factorization.jl b/src/adjoint_factorization.jl index 84306f4f5..325a18a21 100644 --- a/src/adjoint_factorization.jl +++ b/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 diff --git a/src/blocked_lufact.jl b/src/blocked_lufact.jl index ef499e063..470c9a9d8 100644 --- a/src/blocked_lufact.jl +++ b/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 && diff --git a/src/common.jl b/src/common.jl index ea9853c63..ec0a6d911 100644 --- a/src/common.jl +++ b/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 diff --git a/src/default.jl b/src/default.jl index 8bdb9ac6a..c70887f77 100644 --- a/src/default.jl +++ 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] diff --git a/src/factorization.jl b/src/factorization.jl index f94bc82df..42ae3b15f 100644 --- a/src/factorization.jl +++ b/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,18 +693,19 @@ 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( @@ -649,7 +713,9 @@ function init_cacheval( 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)) diff --git a/test/Core/genericlu_naive_ldiv.jl b/test/Core/genericlu_naive_ldiv.jl index 80cd23bd2..9c7241ac4 100644 --- a/test/Core/genericlu_naive_ldiv.jl +++ b/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) diff --git a/test/qa/allocations.jl b/test/qa/allocations.jl index a4433cde3..2132fcd34 100644 --- a/test/qa/allocations.jl +++ 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]