From 2c004a26bb3ad464ac24c432ee6d0be56ef2cbdb Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 09:40:47 +0000 Subject: [PATCH 1/6] Remove ipiv allocation from GenericLUFactorization --- src/LinearSolve.jl | 23 +------------- src/factorization.jl | 70 ++++++++++++++++++++++++++++++++++++++----- src/generic_lufact.jl | 62 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 126 insertions(+), 29 deletions(-) create mode 100644 src/generic_lufact.jl diff --git a/src/LinearSolve.jl b/src/LinearSolve.jl index 63a58a0c3..c7680903b 100644 --- a/src/LinearSolve.jl +++ b/src/LinearSolve.jl @@ -140,6 +140,7 @@ end const BLASELTYPES = Union{Float32, Float64, ComplexF32, ComplexF64} +include("generic_lufact.jl") include("common.jl") include("extension_algs.jl") include("factorization.jl") @@ -171,28 +172,6 @@ end @inline _notsuccessful(F) = hasmethod(LinearAlgebra.issuccess, (typeof(F),)) ? !LinearAlgebra.issuccess(F) : false -@generated function SciMLBase.solve!(cache::LinearCache, alg::AbstractFactorization; - kwargs...) - quote - if cache.isfresh - fact = do_factorization(alg, cache.A, cache.b, cache.u) - cache.cacheval = fact - - # If factorization was not successful, return failure. Don't reset `isfresh` - if _notsuccessful(fact) - return SciMLBase.build_linear_solution( - alg, cache.u, nothing, cache; retcode = ReturnCode.Failure) - end - - cache.isfresh = false - end - - y = _ldiv!(cache.u, @get_cacheval(cache, $(Meta.quot(defaultalg_symbol(alg)))), - cache.b) - return SciMLBase.build_linear_solution(alg, y, nothing, cache; retcode = ReturnCode.Success) - end -end - # Solver Specific Traits ## Needs Square Matrix """ diff --git a/src/factorization.jl b/src/factorization.jl index 84ac5a41d..f5b3e2c5a 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -1,3 +1,25 @@ +@generated function SciMLBase.solve!(cache::LinearCache, alg::AbstractFactorization; + kwargs...) + quote + if cache.isfresh + fact = do_factorization(alg, cache.A, cache.b, cache.u) + cache.cacheval = fact + + # If factorization was not successful, return failure. Don't reset `isfresh` + if _notsuccessful(fact) + return SciMLBase.build_linear_solution( + alg, cache.u, nothing, cache; retcode = ReturnCode.Failure) + end + + cache.isfresh = false + end + + y = _ldiv!(cache.u, @get_cacheval(cache, $(Meta.quot(defaultalg_symbol(alg)))), + cache.b) + return SciMLBase.build_linear_solution(alg, y, nothing, cache; retcode = ReturnCode.Success) + end +end + macro get_cacheval(cache, algsym) quote if $(esc(cache)).alg isa DefaultLinearSolver @@ -8,6 +30,8 @@ macro get_cacheval(cache, algsym) end end +const PREALLOCATED_IPIV = Vector{LinearAlgebra.BlasInt}(undef, 0) + _ldiv!(x, A, b) = ldiv!(x, A, b) _ldiv!(x, A, b::SVector) = (x .= A \ b) @@ -41,8 +65,7 @@ function LinearSolve.init_cacheval( alg::RFLUFactorization, A::Matrix{Float64}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) - ipiv = Vector{LinearAlgebra.BlasInt}(undef, 0) - PREALLOCATED_LU, ipiv + PREALLOCATED_LU, PREALLOCATED_IPIV end function LinearSolve.init_cacheval(alg::RFLUFactorization, @@ -144,14 +167,47 @@ function do_factorization(alg::LUFactorization, A, b, u) return fact end -function do_factorization(alg::GenericLUFactorization, A, b, u) +function init_cacheval( + alg::GenericLUFactorization, A, b, u, Pl, Pr, + maxiters::Int, abstol, reltol, verbose::Bool, + assumptions::OperatorAssumptions) + ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...)) + ArrayInterface.lu_instance(convert(AbstractMatrix, A)), ipiv +end + +function init_cacheval( + alg::GenericLUFactorization, A::Matrix{Float64}, b, u, Pl, Pr, + maxiters::Int, abstol, reltol, verbose::Bool, + assumptions::OperatorAssumptions) + PREALLOCATED_LU, PREALLOCATED_IPIV +end + +function SciMLBase.solve!(cache::LinearSolve.LinearCache, alg::GenericLUFactorization; + kwargs...) + A = cache.A A = convert(AbstractMatrix, A) - fact = LinearAlgebra.generic_lufact!(A, alg.pivot, check = false) - return fact + 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; check = false, ipiv) + cache.cacheval = (fact, ipiv) + + if !LinearAlgebra.issuccess(fact) + return SciMLBase.build_linear_solution( + alg, cache.u, nothing, cache; retcode = ReturnCode.Failure) + end + + cache.isfresh = false + end + y = ldiv!(cache.u, LinearSolve.@get_cacheval(cache, :GenericLUFactorization)[1], cache.b) + SciMLBase.build_linear_solution(alg, y, nothing, cache) end function init_cacheval( - alg::Union{LUFactorization, GenericLUFactorization}, A, b, u, Pl, Pr, + alg::LUFactorization, A, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) ArrayInterface.lu_instance(convert(AbstractMatrix, A)) @@ -171,7 +227,7 @@ end const PREALLOCATED_LU = ArrayInterface.lu_instance(rand(1, 1)) -function init_cacheval(alg::Union{LUFactorization, GenericLUFactorization}, +function init_cacheval(alg::LUFactorization, A::Matrix{Float64}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) diff --git a/src/generic_lufact.jl b/src/generic_lufact.jl new file mode 100644 index 000000000..5fcc06e8b --- /dev/null +++ b/src/generic_lufact.jl @@ -0,0 +1,62 @@ +# From LinearAlgebra.lu.jl +# Modified to be non-allocating +function generic_lufact!(A::AbstractMatrix{T}, pivot::Union{RowMaximum,NoPivot,RowNonZero} = lupivottype(T); + check::Bool = true, ipiv = Vector{BlasInt}(undef, min(size(A)))) where {T} + check && LinearAlgebra.LAPACK.chkfinite(A) + # Extract values + m, n = size(A) + minmn = min(m,n) + + # Initialize variables + info = 0 + + @inbounds begin + for k = 1:minmn + # find index max + kp = k + if pivot === LinearAlgebra.RowMaximum() && k < m + amax = abs(A[k, k]) + for i = k+1:m + absi = abs(A[i,k]) + if absi > amax + kp = i + amax = absi + end + end + elseif pivot === LinearAlgebra.RowNonZero() + for i = k:m + if !iszero(A[i,k]) + kp = i + break + end + end + end + ipiv[k] = kp + if !iszero(A[kp,k]) + if k != kp + # Interchange + for i = 1:n + tmp = A[k,i] + A[k,i] = A[kp,i] + A[kp,i] = tmp + end + end + # Scale first column + Akkinv = inv(A[k,k]) + for i = k+1:m + A[i,k] *= Akkinv + end + elseif info == 0 + info = k + end + # Update the rest + for j = k+1:n + for i = k+1:m + A[i,j] -= A[i,k]*A[k,j] + end + end + end + end + check && LinearAlgebra.checknonsingular(info, pivot) + return LinearAlgebra.LU{T,typeof(A),typeof(ipiv)}(A, ipiv, convert(LinearAlgebra.BlasInt, info)) +end \ No newline at end of file From 7332ca1b1401206af4b3116ec970e2f98d2ee531 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 11:18:55 +0000 Subject: [PATCH 2/6] fix ambiguity --- ext/LinearSolveSparseArraysExt.jl | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/ext/LinearSolveSparseArraysExt.jl b/ext/LinearSolveSparseArraysExt.jl index 9e1a632e0..d04e8117f 100644 --- a/ext/LinearSolveSparseArraysExt.jl +++ b/ext/LinearSolveSparseArraysExt.jl @@ -64,7 +64,15 @@ const PREALLOCATED_UMFPACK = SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(0, 0 Int[], Float64[])) function LinearSolve.init_cacheval( - alg::Union{LUFactorization, GenericLUFactorization}, A::AbstractSparseArray{<:Number, <:Integer}, b, u, + alg::LUFactorization, A::AbstractSparseArray{<:Number, <:Integer}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Bool, assumptions::OperatorAssumptions) + nothing +end + +function LinearSolve.init_cacheval( + alg::GenericLUFactorization, A::AbstractSparseArray{<:Number, <:Integer}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) @@ -80,7 +88,7 @@ function LinearSolve.init_cacheval( end function LinearSolve.init_cacheval( - alg::Union{LUFactorization, GenericLUFactorization}, A::AbstractSparseArray{Float64, Int64}, b, u, + alg::LUFactorization, A::AbstractSparseArray{Float64, Int64}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) @@ -88,7 +96,7 @@ function LinearSolve.init_cacheval( end function LinearSolve.init_cacheval( - alg::Union{LUFactorization, GenericLUFactorization}, A::AbstractSparseArray{T, Int64}, b, u, + alg::LUFactorization, A::AbstractSparseArray{T, Int64}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) where {T<:BLASELTYPES} @@ -96,7 +104,7 @@ function LinearSolve.init_cacheval( end function LinearSolve.init_cacheval( - alg::Union{LUFactorization, GenericLUFactorization}, A::AbstractSparseArray{T, Int32}, b, u, + alg::LUFactorization, A::AbstractSparseArray{T, Int32}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) where {T<:BLASELTYPES} From d51e5d09e32efd0c07425b942988df22614f2103 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 11:39:19 +0000 Subject: [PATCH 3/6] fix ambiguity --- src/factorization.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/factorization.jl b/src/factorization.jl index f5b3e2c5a..e670d185d 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -234,7 +234,14 @@ function init_cacheval(alg::LUFactorization, PREALLOCATED_LU end -function init_cacheval(alg::Union{LUFactorization, GenericLUFactorization}, +function init_cacheval(alg::LUFactorization, + A::AbstractSciMLOperator, b, u, Pl, Pr, + maxiters::Int, abstol, reltol, verbose::Bool, + assumptions::OperatorAssumptions) + nothing +end + +function init_cacheval(alg::GenericLUFactorization, A::AbstractSciMLOperator, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) From 7e8c2fc13a6dd09df70879b176884d34039155be Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 14:14:13 +0000 Subject: [PATCH 4/6] fix kwarg definition --- src/generic_lufact.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/generic_lufact.jl b/src/generic_lufact.jl index 5fcc06e8b..45b0ea2b0 100644 --- a/src/generic_lufact.jl +++ b/src/generic_lufact.jl @@ -1,7 +1,7 @@ # From LinearAlgebra.lu.jl # Modified to be non-allocating function generic_lufact!(A::AbstractMatrix{T}, pivot::Union{RowMaximum,NoPivot,RowNonZero} = lupivottype(T); - check::Bool = true, ipiv = Vector{BlasInt}(undef, min(size(A)))) where {T} + check::Bool = true, ipiv = Vector{BlasInt}(undef, min(size(A)...))) where {T} check && LinearAlgebra.LAPACK.chkfinite(A) # Extract values m, n = size(A) From d315d13482e3fb0ccfaa23b4292f741777936fa0 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 14:35:56 +0000 Subject: [PATCH 5/6] match the v1.13 format and setup different versions for generic_lufact --- src/factorization.jl | 3 +- src/generic_lufact.jl | 164 ++++++++++++++++++++++++++++++------------ 2 files changed, 120 insertions(+), 47 deletions(-) diff --git a/src/factorization.jl b/src/factorization.jl index e670d185d..a0616b6aa 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -192,7 +192,7 @@ function SciMLBase.solve!(cache::LinearSolve.LinearCache, alg::GenericLUFactoriz if length(ipiv) != min(size(A)...) ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...)) end - fact = generic_lufact!(A, alg.pivot; check = false, ipiv) + fact = generic_lufact!(A, alg.pivot, ipiv; check = false) cache.cacheval = (fact, ipiv) if !LinearAlgebra.issuccess(fact) @@ -221,6 +221,7 @@ function init_cacheval(alg::Union{LUFactorization, GenericLUFactorization}, return lu(A; check = false) else A isa GPUArraysCore.AnyGPUArray && return nothing + ipiv = Vector{LinearAlgebra.BlasInt}(undef, 0) return LinearAlgebra.generic_lufact!(copy(A), alg.pivot; check = false) end end diff --git a/src/generic_lufact.jl b/src/generic_lufact.jl index 45b0ea2b0..36478376e 100644 --- a/src/generic_lufact.jl +++ b/src/generic_lufact.jl @@ -1,62 +1,134 @@ # From LinearAlgebra.lu.jl # Modified to be non-allocating -function generic_lufact!(A::AbstractMatrix{T}, pivot::Union{RowMaximum,NoPivot,RowNonZero} = lupivottype(T); - check::Bool = true, ipiv = Vector{BlasInt}(undef, min(size(A)...))) where {T} - check && LinearAlgebra.LAPACK.chkfinite(A) - # Extract values - m, n = size(A) - minmn = min(m,n) +@static if VERSION < v"1.11" + function generic_lufact!(A::AbstractMatrix{T}, pivot::Union{RowMaximum,NoPivot,RowNonZero} = LinearAlgebra.lupivottype(T), + ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...)); + check::Bool = true, allowsingular::Bool = false) where {T} + check && LinearAlgebra.LAPACK.chkfinite(A) + # Extract values + m, n = size(A) + minmn = min(m,n) - # Initialize variables - info = 0 - - @inbounds begin - for k = 1:minmn - # find index max - kp = k - if pivot === LinearAlgebra.RowMaximum() && k < m - amax = abs(A[k, k]) - for i = k+1:m - absi = abs(A[i,k]) - if absi > amax - kp = i - amax = absi + # Initialize variables + info = 0 + + @inbounds begin + for k = 1:minmn + # find index max + kp = k + if pivot === LinearAlgebra.RowMaximum() && k < m + amax = abs(A[k, k]) + for i = k+1:m + absi = abs(A[i,k]) + if absi > amax + kp = i + amax = absi + end end + elseif pivot === LinearAlgebra.RowNonZero() + for i = k:m + if !iszero(A[i,k]) + kp = i + break + end + end + end + ipiv[k] = kp + if !iszero(A[kp,k]) + if k != kp + # Interchange + for i = 1:n + tmp = A[k,i] + A[k,i] = A[kp,i] + A[kp,i] = tmp + end + end + # Scale first column + Akkinv = inv(A[k,k]) + for i = k+1:m + A[i,k] *= Akkinv + end + elseif info == 0 + info = k end - elseif pivot === LinearAlgebra.RowNonZero() - for i = k:m - if !iszero(A[i,k]) - kp = i - break + # Update the rest + for j = k+1:n + for i = k+1:m + A[i,j] -= A[i,k]*A[k,j] end end end - ipiv[k] = kp - if !iszero(A[kp,k]) - if k != kp - # Interchange - for i = 1:n - tmp = A[k,i] - A[k,i] = A[kp,i] - A[kp,i] = tmp + end + check && LinearAlgebra.checknonsingular(info, pivot) + return LinearAlgebra.LU{T,typeof(A),typeof(ipiv)}(A, ipiv, convert(LinearAlgebra.BlasInt, info)) + end +elseif VERSION < v"1.13" + function generic_lufact!(A::AbstractMatrix{T}, pivot::Union{RowMaximum,NoPivot,RowNonZero} = LinearAlgebra.lupivottype(T), + ipiv = Vector{LinearAlgebra.BlasInt}(undef, min(size(A)...)); + check::Bool = true, allowsingular::Bool = false) where {T} + check && LAPACK.chkfinite(A) + # Extract values + m, n = size(A) + minmn = min(m,n) + + # Initialize variables + info = 0 + + @inbounds begin + for k = 1:minmn + # find index max + kp = k + if pivot === LinearAlgebra.RowMaximum() && k < m + amax = abs(A[k, k]) + for i = k+1:m + absi = abs(A[i,k]) + if absi > amax + kp = i + amax = absi + end + end + elseif pivot === LinearAlgebra.RowNonZero() + for i = k:m + if !iszero(A[i,k]) + kp = i + break + end end end - # Scale first column - Akkinv = inv(A[k,k]) - for i = k+1:m - A[i,k] *= Akkinv + ipiv[k] = kp + if !iszero(A[kp,k]) + if k != kp + # Interchange + for i = 1:n + tmp = A[k,i] + A[k,i] = A[kp,i] + A[kp,i] = tmp + end + end + # Scale first column + Akkinv = inv(A[k,k]) + for i = k+1:m + A[i,k] *= Akkinv + end + elseif info == 0 + info = k end - elseif info == 0 - info = k - end - # Update the rest - for j = k+1:n - for i = k+1:m - A[i,j] -= A[i,k]*A[k,j] + # Update the rest + for j = k+1:n + for i = k+1:m + A[i,j] -= A[i,k]*A[k,j] + end end end end + if pivot === LinearAlgebra.NoPivot() + # Use a negative value to distinguish a failed factorization (zero in pivot + # position during unpivoted LU) from a valid but rank-deficient factorization + info = -info + end + check && LinearAlgebra._check_lu_success(info, allowsingular) + return LinearAlgebra.LU{T,typeof(A),typeof(ipiv)}(A, ipiv, convert(LinearAlgebra.BlasInt, info)) end - check && LinearAlgebra.checknonsingular(info, pivot) - return LinearAlgebra.LU{T,typeof(A),typeof(ipiv)}(A, ipiv, convert(LinearAlgebra.BlasInt, info)) +else + generic_lufact!(args...; kwargs...) = LinearAlgebra.generic_lufact!(args...; kwargs...) end \ No newline at end of file From 62cdcbfb71d935edeb12b1af14acc219141d20dd Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 26 May 2025 14:55:26 +0000 Subject: [PATCH 6/6] fix GPU tests --- src/factorization.jl | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/src/factorization.jl b/src/factorization.jl index a0616b6aa..8d03a5d2c 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -213,17 +213,20 @@ function init_cacheval( ArrayInterface.lu_instance(convert(AbstractMatrix, A)) end -function init_cacheval(alg::Union{LUFactorization, GenericLUFactorization}, +function init_cacheval(alg::LUFactorization, A::Union{<:Adjoint, <:Transpose}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) error_no_cudss_lu(A) - if alg isa LUFactorization - return lu(A; check = false) - else - A isa GPUArraysCore.AnyGPUArray && return nothing - ipiv = Vector{LinearAlgebra.BlasInt}(undef, 0) - return LinearAlgebra.generic_lufact!(copy(A), alg.pivot; check = false) - end + return lu(A; check = false) +end + +function init_cacheval(alg::GenericLUFactorization, + A::Union{<:Adjoint, <:Transpose}, b, u, Pl, Pr, maxiters::Int, abstol, reltol, + verbose::Bool, assumptions::OperatorAssumptions) + error_no_cudss_lu(A) + A isa GPUArraysCore.AnyGPUArray && return nothing + ipiv = Vector{LinearAlgebra.BlasInt}(undef, 0) + return LinearAlgebra.generic_lufact!(copy(A), alg.pivot; check = false), ipiv end const PREALLOCATED_LU = ArrayInterface.lu_instance(rand(1, 1))