diff --git a/ext/LinearSolveSparseArraysExt.jl b/ext/LinearSolveSparseArraysExt.jl index 70f17047d..44f3f7472 100644 --- a/ext/LinearSolveSparseArraysExt.jl +++ b/ext/LinearSolveSparseArraysExt.jl @@ -53,15 +53,15 @@ function LinearSolve.handle_sparsematrixcsc_lu(A::AbstractSparseMatrixCSC) end @static if Base.USE_GPL_LIBS -function LinearSolve.defaultalg( - A::Symmetric{<:BLASELTYPES, <:SparseMatrixCSC}, b, ::OperatorAssumptions{Bool}) - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.CHOLMODFactorization) -end + function LinearSolve.defaultalg( + A::Symmetric{<:BLASELTYPES, <:SparseMatrixCSC}, b, ::OperatorAssumptions{Bool}) + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.CHOLMODFactorization) + end else -function LinearSolve.defaultalg( - A::Symmetric{<:BLASELTYPES, <:SparseMatrixCSC}, b, ::OperatorAssumptions{Bool}) - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.CholeskyFactorization) -end + function LinearSolve.defaultalg( + A::Symmetric{<:BLASELTYPES, <:SparseMatrixCSC}, b, ::OperatorAssumptions{Bool}) + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.CholeskyFactorization) + end end # @static if Base.USE_GPL_LIBS function LinearSolve.defaultalg(A::AbstractSparseMatrixCSC{Tv, Ti}, b, @@ -88,8 +88,8 @@ function LinearSolve.init_cacheval(alg::GenericFactorization, end @static if Base.USE_GPL_LIBS -const PREALLOCATED_UMFPACK = SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(0, 0, [1], - Int[], Float64[])) + const PREALLOCATED_UMFPACK = SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(0, 0, [1], + Int[], Float64[])) end # @static if Base.USE_GPL_LIBS function LinearSolve.init_cacheval( @@ -117,37 +117,39 @@ function LinearSolve.init_cacheval( end @static if Base.USE_GPL_LIBS -function LinearSolve.init_cacheval( - alg::LUFactorization, A::AbstractSparseArray{Float64, Int64}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) - PREALLOCATED_UMFPACK -end -function LinearSolve.init_cacheval( - alg::LUFactorization, A::AbstractSparseArray{T, Int64}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: BLASELTYPES} - if LinearSolve.is_cusparse(A) - LinearSolve.cudss_loaded(A) ? ArrayInterface.lu_instance(A) : nothing - else - SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int64}( - zero(Int64), zero(Int64), [Int64(1)], Int64[], T[])) + function LinearSolve.init_cacheval( + alg::LUFactorization, A::AbstractSparseArray{Float64, Int64}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + PREALLOCATED_UMFPACK end -end -function LinearSolve.init_cacheval( - alg::LUFactorization, A::AbstractSparseArray{T, Int32}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: BLASELTYPES} - if LinearSolve.is_cusparse(A) - LinearSolve.cudss_loaded(A) ? ArrayInterface.lu_instance(A) : nothing - else - SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int32}( - zero(Int32), zero(Int32), [Int32(1)], Int32[], T[])) + function LinearSolve.init_cacheval( + alg::LUFactorization, A::AbstractSparseArray{T, Int64}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: + BLASELTYPES} + if LinearSolve.is_cusparse(A) + LinearSolve.cudss_loaded(A) ? ArrayInterface.lu_instance(A) : nothing + else + SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int64}( + zero(Int64), zero(Int64), [Int64(1)], Int64[], T[])) + end + end + function LinearSolve.init_cacheval( + alg::LUFactorization, A::AbstractSparseArray{T, Int32}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: + BLASELTYPES} + if LinearSolve.is_cusparse(A) + LinearSolve.cudss_loaded(A) ? ArrayInterface.lu_instance(A) : nothing + else + SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int32}( + zero(Int32), zero(Int32), [Int32(1)], Int32[], T[])) + end end -end end # @static if Base.USE_GPL_LIBS function LinearSolve.init_cacheval( @@ -167,76 +169,79 @@ function LinearSolve.init_cacheval( end @static if Base.USE_GPL_LIBS -function LinearSolve.init_cacheval( - alg::UMFPACKFactorization, A::AbstractSparseArray{Float64, Int}, b, u, Pl, Pr, - maxiters::Int, abstol, - reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + function LinearSolve.init_cacheval( + alg::UMFPACKFactorization, A::AbstractSparseArray{Float64, Int}, b, u, Pl, Pr, + maxiters::Int, abstol, + reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) PREALLOCATED_UMFPACK -end + end -function LinearSolve.init_cacheval( - alg::UMFPACKFactorization, A::AbstractSparseArray{T, Int64}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: BLASELTYPES} - SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int64}( - zero(Int64), zero(Int64), [Int64(1)], Int64[], T[])) -end + function LinearSolve.init_cacheval( + alg::UMFPACKFactorization, A::AbstractSparseArray{T, Int64}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: + BLASELTYPES} + SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int64}( + zero(Int64), zero(Int64), [Int64(1)], Int64[], T[])) + end -function LinearSolve.init_cacheval( - alg::UMFPACKFactorization, A::AbstractSparseArray{T, Int32}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: BLASELTYPES} - SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int32}( - zero(Int32), zero(Int32), [Int32(1)], Int32[], T[])) -end + function LinearSolve.init_cacheval( + alg::UMFPACKFactorization, A::AbstractSparseArray{T, Int32}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: + BLASELTYPES} + SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int32}( + zero(Int32), zero(Int32), [Int32(1)], Int32[], T[])) + end -function SciMLBase.solve!( - cache::LinearSolve.LinearCache, alg::UMFPACKFactorization; kwargs...) - A = cache.A - A = convert(AbstractMatrix, A) - if cache.isfresh - cacheval = LinearSolve.@get_cacheval(cache, :UMFPACKFactorization) - if alg.reuse_symbolic - # Caches the symbolic factorization: https://github.com/JuliaLang/julia/pull/33738 - if length(cacheval.nzval) != length(nonzeros(A)) || alg.check_pattern && pattern_changed(cacheval, A) + function SciMLBase.solve!( + cache::LinearSolve.LinearCache, alg::UMFPACKFactorization; kwargs...) + A = cache.A + A = convert(AbstractMatrix, A) + if cache.isfresh + cacheval = LinearSolve.@get_cacheval(cache, :UMFPACKFactorization) + if alg.reuse_symbolic + # Caches the symbolic factorization: https://github.com/JuliaLang/julia/pull/33738 + if length(cacheval.nzval) != length(nonzeros(A)) || + alg.check_pattern && pattern_changed(cacheval, A) + fact = lu( + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), + nonzeros(A)), + check = false) + else + fact = lu!(cacheval, + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), + nonzeros(A)), check = false) + end + else fact = lu( - SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), - nonzeros(A)), + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)), check = false) - else - fact = lu!(cacheval, - SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), - nonzeros(A)), check = false) end - else - fact = lu( - SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)), - check = false) + cache.cacheval = fact + cache.isfresh = false end - cache.cacheval = fact - cache.isfresh = false - end - F = LinearSolve.@get_cacheval(cache, :UMFPACKFactorization) - if F.status == UMFPACK_OK - y = ldiv!(cache.u, F, cache.b) - SciMLBase.build_linear_solution( - alg, y, nothing, cache; retcode = ReturnCode.Success) - else - @SciMLMessage("Solver failed", cache.verbose, :solver_failure) - SciMLBase.build_linear_solution( - alg, cache.u, nothing, cache; retcode = ReturnCode.Infeasible) + F = LinearSolve.@get_cacheval(cache, :UMFPACKFactorization) + if F.status == UMFPACK_OK + y = ldiv!(cache.u, F, cache.b) + SciMLBase.build_linear_solution( + alg, y, nothing, cache; retcode = ReturnCode.Success) + else + @SciMLMessage("Solver failed", cache.verbose, :solver_failure) + SciMLBase.build_linear_solution( + alg, cache.u, nothing, cache; retcode = ReturnCode.Infeasible) + end end -end else -function SciMLBase.solve!( - cache::LinearSolve.LinearCache, alg::UMFPACKFactorization; kwargs...) - error("UMFPACKFactorization requires GPL libraries (UMFPACK). Rebuild Julia with USE_GPL_LIBS=1 or use an alternative algorithm like SparspakFactorization") -end + function SciMLBase.solve!( + cache::LinearSolve.LinearCache, alg::UMFPACKFactorization; kwargs...) + error("UMFPACKFactorization requires GPL libraries (UMFPACK). Rebuild Julia with USE_GPL_LIBS=1 or use an alternative algorithm like SparspakFactorization") + end end # @static if Base.USE_GPL_LIBS function LinearSolve.init_cacheval( @@ -312,7 +317,6 @@ function LinearSolve.init_cacheval( end end - function LinearSolve.init_cacheval( alg::NormalCholeskyFactorization, A::AbstractSciMLOperator, b, u, Pl, Pr, maxiters::Int, abstol, reltol, @@ -331,7 +335,8 @@ function SciMLBase.solve!(cache::LinearSolve.LinearCache, alg::KLUFactorization; if cache.isfresh cacheval = LinearSolve.@get_cacheval(cache, :KLUFactorization) if alg.reuse_symbolic - if length(cacheval.nzval) != length(nonzeros(A)) || alg.check_pattern && pattern_changed(cacheval, A) + if length(cacheval.nzval) != length(nonzeros(A)) || + alg.check_pattern && pattern_changed(cacheval, A) fact = KLU.klu( SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)), @@ -369,25 +374,25 @@ function LinearSolve.init_cacheval(alg::CHOLMODFactorization, end @static if Base.USE_GPL_LIBS -const PREALLOCATED_CHOLMOD = cholesky(sparse(reshape([1.0], 1, 1))) - -function LinearSolve.init_cacheval(alg::CHOLMODFactorization, - A::Union{SparseMatrixCSC{T, Int}, Symmetric{T, SparseMatrixCSC{T, Int}}}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: - Float64} - PREALLOCATED_CHOLMOD -end + const PREALLOCATED_CHOLMOD = cholesky(sparse(reshape([1.0], 1, 1))) + + function LinearSolve.init_cacheval(alg::CHOLMODFactorization, + A::Union{SparseMatrixCSC{T, Int}, Symmetric{T, SparseMatrixCSC{T, Int}}}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: + Float64} + PREALLOCATED_CHOLMOD + end -function LinearSolve.init_cacheval(alg::CHOLMODFactorization, - A::Union{SparseMatrixCSC{T, Int}, Symmetric{T, SparseMatrixCSC{T, Int}}}, b, u, - Pl, Pr, - maxiters::Int, abstol, reltol, - verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: - BLASELTYPES} - cholesky(sparse(reshape([one(T)], 1, 1))) -end + function LinearSolve.init_cacheval(alg::CHOLMODFactorization, + A::Union{SparseMatrixCSC{T, Int}, Symmetric{T, SparseMatrixCSC{T, Int}}}, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) where {T <: + BLASELTYPES} + cholesky(sparse(reshape([one(T)], 1, 1))) + end end # @static if Base.USE_GPL_LIBS function LinearSolve.init_cacheval(alg::NormalCholeskyFactorization, @@ -399,6 +404,11 @@ function LinearSolve.init_cacheval(alg::NormalCholeskyFactorization, nothing elseif LinearSolve.is_cusparse_csr(A) && !LinearSolve.cudss_loaded(A) nothing + elseif A isa LinearSolve.GPUArraysCore.AnyGPUArray && !assumptions.issq + # GPU arrays need special handling for non-square matrices to avoid + # DimensionMismatch errors during DefaultLinearSolver cache initialization + # See https://github.com/SciML/NonlinearSolve.jl/issues/746 + nothing else ArrayInterface.cholesky_instance(convert(AbstractMatrix, A)) end @@ -430,24 +440,24 @@ end @static if Base.USE_GPL_LIBS # SPQR and CHOLMOD Factor support -function LinearSolve._ldiv!(x::Vector, - A::Union{SparseArrays.SPQR.QRSparse, SparseArrays.CHOLMOD.Factor}, b::Vector) - x .= A \ b -end -function LinearSolve._ldiv!(x::AbstractVector, - A::Union{SparseArrays.SPQR.QRSparse, SparseArrays.CHOLMOD.Factor}, b::AbstractVector) - x .= A \ b -end -function LinearSolve._ldiv!(::SVector, - A::Union{SparseArrays.CHOLMOD.Factor, SparseArrays.SPQR.QRSparse}, - b::AbstractVector) - (A \ b) -end -function LinearSolve._ldiv!(::SVector, - A::Union{SparseArrays.CHOLMOD.Factor, SparseArrays.SPQR.QRSparse}, - b::SVector) - (A \ b) -end + function LinearSolve._ldiv!(x::Vector, + A::Union{SparseArrays.SPQR.QRSparse, SparseArrays.CHOLMOD.Factor}, b::Vector) + x .= A \ b + end + function LinearSolve._ldiv!(x::AbstractVector, + A::Union{SparseArrays.SPQR.QRSparse, SparseArrays.CHOLMOD.Factor}, b::AbstractVector) + x .= A \ b + end + function LinearSolve._ldiv!(::SVector, + A::Union{SparseArrays.CHOLMOD.Factor, SparseArrays.SPQR.QRSparse}, + b::AbstractVector) + (A \ b) + end + function LinearSolve._ldiv!(::SVector, + A::Union{SparseArrays.CHOLMOD.Factor, SparseArrays.SPQR.QRSparse}, + b::SVector) + (A \ b) + end end # @static if Base.USE_GPL_LIBS function LinearSolve.pattern_changed(fact::Nothing, A::SparseArrays.SparseMatrixCSC) @@ -471,29 +481,29 @@ function LinearSolve.pattern_changed(fact, A::SparseArrays.SparseMatrixCSC) end @static if Base.USE_GPL_LIBS -function LinearSolve.defaultalg( - A::AbstractSparseMatrixCSC{<:Union{Float64, ComplexF64}, Ti}, b, - assump::OperatorAssumptions{Bool}) where {Ti} - if assump.issq - if length(b) <= 10_000 && length(nonzeros(A)) / length(A) < 2e-4 - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KLUFactorization) + function LinearSolve.defaultalg( + A::AbstractSparseMatrixCSC{<:Union{Float64, ComplexF64}, Ti}, b, + assump::OperatorAssumptions{Bool}) where {Ti} + if assump.issq + if length(b) <= 10_000 && length(nonzeros(A)) / length(A) < 2e-4 + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KLUFactorization) + else + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.UMFPACKFactorization) + end else - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.UMFPACKFactorization) + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.QRFactorization) end - else - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.QRFactorization) end -end else -function LinearSolve.defaultalg( - A::AbstractSparseMatrixCSC{<:Union{Float64, ComplexF64}, Ti}, b, - assump::OperatorAssumptions{Bool}) where {Ti} - if assump.issq - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KLUFactorization) - elseif !assump.issq - LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.QRFactorization) + function LinearSolve.defaultalg( + A::AbstractSparseMatrixCSC{<:Union{Float64, ComplexF64}, Ti}, b, + assump::OperatorAssumptions{Bool}) where {Ti} + if assump.issq + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KLUFactorization) + elseif !assump.issq + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.QRFactorization) + end end -end end # @static if Base.USE_GPL_LIBS # SPQR Handling diff --git a/src/factorization.jl b/src/factorization.jl index ed7aaa723..3056a2cf5 100644 --- a/src/factorization.jl +++ b/src/factorization.jl @@ -317,13 +317,15 @@ end const PREALLOCATED_QR_ColumnNorm = ArrayInterface.qr_instance(rand(1, 1), ColumnNorm()) function init_cacheval(alg::QRFactorization{ColumnNorm}, A::Matrix{Float64}, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) return PREALLOCATED_QR_ColumnNorm end function init_cacheval( alg::QRFactorization, A::Union{<:Adjoint, <:Transpose}, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) A isa GPUArraysCore.AnyGPUArray && return qr(A) return qr(A, alg.pivot) end @@ -331,7 +333,8 @@ end const PREALLOCATED_QR_NoPivot = ArrayInterface.qr_instance(rand(1, 1)) function init_cacheval(alg::QRFactorization{NoPivot}, A::Matrix{Float64}, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) return PREALLOCATED_QR_NoPivot end @@ -388,13 +391,18 @@ function init_cacheval(alg::CholeskyFactorization, A::SMatrix{S1, S2}, b, u, Pl, end function init_cacheval(alg::CholeskyFactorization, A::GPUArraysCore.AnyGPUArray, b, u, Pl, - Pr, maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) - cholesky(A; check = false) + Pr, maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) + # Cholesky requires square matrices - return nothing for non-square to avoid errors + # during DefaultLinearSolver cache initialization + # See https://github.com/SciML/NonlinearSolve.jl/issues/746 + assumptions.issq ? cholesky(A; check = false) : nothing end function init_cacheval( alg::CholeskyFactorization, A::AbstractArray{<:BLASELTYPES}, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) if LinearSolve.is_cusparse_csc(A) nothing elseif LinearSolve.is_cusparse_csr(A) && !LinearSolve.cudss_loaded(A) @@ -1012,7 +1020,8 @@ const PREALLOCATED_NORMALCHOLESKY_SYMMETRIC = ArrayInterface.cholesky_instance( Symmetric(rand(1, 1)), NoPivot()) function init_cacheval(alg::NormalCholeskyFactorization, A::Matrix{Float64}, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) return PREALLOCATED_NORMALCHOLESKY_SYMMETRIC end @@ -1164,7 +1173,8 @@ function init_cacheval(alg::SparspakFactorization, end function init_cacheval(::SparspakFactorization, ::StaticArray, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) nothing end @@ -1190,9 +1200,8 @@ struct CliqueTreesFactorization{A, S} <: AbstractSparseFactorization alg::A = nothing, snd::S = nothing, reuse_symbolic = true, - throwerror = true, - ) where {A, S} - + throwerror = true + ) where {A, S} ext = Base.get_extension(@__MODULE__, :LinearSolveCliqueTreesExt) if throwerror && isnothing(ext) @@ -1203,30 +1212,36 @@ struct CliqueTreesFactorization{A, S} <: AbstractSparseFactorization end end -function init_cacheval(::CliqueTreesFactorization, ::Union{AbstractMatrix, Nothing, AbstractSciMLOperator}, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) +function init_cacheval(::CliqueTreesFactorization, + ::Union{AbstractMatrix, Nothing, AbstractSciMLOperator}, b, u, Pl, Pr, + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) nothing end function init_cacheval(::CliqueTreesFactorization, ::StaticArray, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) nothing end # Fallback init_cacheval for extension-based algorithms when extensions aren't loaded # These return nothing since the actual implementations are in the extensions function init_cacheval(::BLISLUFactorization, A, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) nothing end function init_cacheval(::CudaOffloadLUFactorization, A, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) nothing end function init_cacheval(::MetalLUFactorization, A, b, u, Pl, Pr, - maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions) + maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, + assumptions::OperatorAssumptions) nothing end diff --git a/test/gpu/cuda.jl b/test/gpu/cuda.jl index 8a3edf8d1..24b2a6d90 100644 --- a/test/gpu/cuda.jl +++ b/test/gpu/cuda.jl @@ -75,7 +75,7 @@ x2 = zero(b); prob1 = LinearProblem(A1, b1; u0 = x1) prob2 = LinearProblem(A2, b2; u0 = x2) -cache_kwargs = (;abstol = 1e-8, reltol = 1e-8, maxiter = 30) +cache_kwargs = (; abstol = 1e-8, reltol = 1e-8, maxiter = 30) function test_interface(alg, prob1, prob2) A1 = prob1.A @@ -103,7 +103,8 @@ function test_interface(alg, prob1, prob2) return end -@testset "$alg" for alg in (CudaOffloadLUFactorization(), CudaOffloadQRFactorization(), NormalCholeskyFactorization()) +@testset "$alg" for alg in (CudaOffloadLUFactorization(), CudaOffloadQRFactorization(), + NormalCholeskyFactorization()) test_interface(alg, prob1, prob2) end @@ -171,3 +172,45 @@ if Base.find_package("CUSOLVERRF") !== nothing include("cusolverrf.jl") end end + +# Test for non-square GPU matrices (least squares problems) +# See https://github.com/SciML/NonlinearSolve.jl/issues/746 +@testset "Non-square GPU matrices" begin + # Overdetermined system: more rows than columns (4x2) + A_rect = cu(Float32[1.0 2.0; 3.0 4.0; 5.0 6.0; 7.0 8.0]) + b_rect = cu(Float32[1.0, 2.0, 3.0, 4.0]) + + prob_rect = LinearProblem(A_rect, b_rect) + + # Test that default solver works (should use QRFactorization) + @testset "Default solver for non-square" begin + sol = solve(prob_rect) + @test sol.alg.alg == LinearSolve.DefaultAlgorithmChoice.QRFactorization + # Verify least squares solution + @test norm(A_rect * sol.u - b_rect) < norm(b_rect) # residual should be smaller than b + end + + # Test explicit QRFactorization + @testset "QRFactorization for non-square" begin + sol = solve(prob_rect, QRFactorization()) + @test norm(A_rect * sol.u - b_rect) < norm(b_rect) + end + + # Test NormalCholeskyFactorization (should work via A'*A) + @testset "NormalCholeskyFactorization for non-square" begin + sol = solve(prob_rect, NormalCholeskyFactorization()) + @test norm(A_rect * sol.u - b_rect) < norm(b_rect) + end + + # Underdetermined system: more columns than rows (2x4) + A_under = cu(Float32[1.0 2.0 3.0 4.0; 5.0 6.0 7.0 8.0]) + b_under = cu(Float32[1.0, 2.0]) + + prob_under = LinearProblem(A_under, b_under) + + @testset "Default solver for underdetermined" begin + sol = solve(prob_under) + # Should still work and give a solution + @test norm(A_under * sol.u - b_under) < 1e-4 + end +end