From e97c844ef26fb9483aa1905cd4eae30211eb786e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 22 Apr 2024 16:49:18 -0400 Subject: [PATCH] early failure on bad factorization --- src/LinearSolve.jl | 523 ++++++++++++++++++++++--------------------- src/default.jl | 14 +- test/default_algs.jl | 6 +- 3 files changed, 280 insertions(+), 263 deletions(-) diff --git a/src/LinearSolve.jl b/src/LinearSolve.jl index 57e6d5bd0..f1633d430 100644 --- a/src/LinearSolve.jl +++ b/src/LinearSolve.jl @@ -1,258 +1,265 @@ -module LinearSolve -if isdefined(Base, :Experimental) && - isdefined(Base.Experimental, Symbol("@max_methods")) - @eval Base.Experimental.@max_methods 1 -end - -import PrecompileTools - -PrecompileTools.@recompile_invalidations begin - using ArrayInterface - using RecursiveFactorization - using Base: cache_dependencies, Bool - using LinearAlgebra - using SparseArrays - using SparseArrays: AbstractSparseMatrixCSC, nonzeros, rowvals, getcolptr - using LazyArrays: @~, BroadcastArray - using SciMLBase: AbstractLinearAlgorithm - using SciMLOperators - using SciMLOperators: AbstractSciMLOperator, IdentityOperator - using Setfield - using UnPack - using KLU - using Sparspak - using FastLapackInterface - using DocStringExtensions - using EnumX - using Markdown - using ChainRulesCore - import InteractiveUtils - - import StaticArraysCore: StaticArray, SVector, MVector, SMatrix, MMatrix - - using LinearAlgebra: BlasInt, LU - using LinearAlgebra.LAPACK: require_one_based_indexing, - chkfinite, chkstride1, - @blasfunc, chkargsok - - import GPUArraysCore - import Preferences - import ConcreteStructs: @concrete - - # wrap - import Krylov - using SciMLBase - import Preferences -end - -const CRC = ChainRulesCore - -if Preferences.@load_preference("LoadMKL_JLL", true) - using MKL_jll - const usemkl = MKL_jll.is_available() -else - const usemkl = false -end - -using Reexport -@reexport using SciMLBase -using SciMLBase: _unwrap_val - -abstract type SciMLLinearSolveAlgorithm <: SciMLBase.AbstractLinearAlgorithm end -abstract type AbstractFactorization <: SciMLLinearSolveAlgorithm end -abstract type AbstractKrylovSubspaceMethod <: SciMLLinearSolveAlgorithm end -abstract type AbstractSolveFunction <: SciMLLinearSolveAlgorithm end - -# Traits - -needs_concrete_A(alg::AbstractFactorization) = true -needs_concrete_A(alg::AbstractKrylovSubspaceMethod) = false -needs_concrete_A(alg::AbstractSolveFunction) = false - -# Util -is_underdetermined(x) = false -is_underdetermined(A::AbstractMatrix) = size(A, 1) < size(A, 2) -is_underdetermined(A::AbstractSciMLOperator) = size(A, 1) < size(A, 2) - -_isidentity_struct(A) = false -_isidentity_struct(λ::Number) = isone(λ) -_isidentity_struct(A::UniformScaling) = isone(A.λ) -_isidentity_struct(::SciMLOperators.IdentityOperator) = true - -# Dispatch Friendly way to check if an extension is loaded -__is_extension_loaded(::Val) = false - -# Check if a sparsity pattern has changed -pattern_changed(fact, A) = false - -function _fast_sym_givens! end - -# Code - -const INCLUDE_SPARSE = Preferences.@load_preference("include_sparse", Base.USE_GPL_LIBS) - -EnumX.@enumx DefaultAlgorithmChoice begin - LUFactorization - QRFactorization - DiagonalFactorization - DirectLdiv! - SparspakFactorization - KLUFactorization - UMFPACKFactorization - KrylovJL_GMRES - GenericLUFactorization - RFLUFactorization - LDLtFactorization - BunchKaufmanFactorization - CHOLMODFactorization - SVDFactorization - CholeskyFactorization - NormalCholeskyFactorization - AppleAccelerateLUFactorization - MKLLUFactorization - QRFactorizationPivoted - KrylovJL_CRAIGMR - KrylovJL_LSMR -end - -struct DefaultLinearSolver <: SciMLLinearSolveAlgorithm - alg::DefaultAlgorithmChoice.T -end - -const BLASELTYPES = Union{Float32, Float64, ComplexF32, ComplexF64} - -include("common.jl") -include("factorization.jl") -include("appleaccelerate.jl") -include("mkl.jl") -include("simplelu.jl") -include("simplegmres.jl") -include("iterative_wrappers.jl") -include("preconditioners.jl") -include("solve_function.jl") -include("default.jl") -include("init.jl") -include("extension_algs.jl") -include("adjoint.jl") -include("deprecated.jl") - -@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 - cache.isfresh = false - end - y = _ldiv!(cache.u, @get_cacheval(cache, $(Meta.quot(defaultalg_symbol(alg)))), - cache.b) - - #= - retcode = if LinearAlgebra.issuccess(fact) - SciMLBase.ReturnCode.Success - else - SciMLBase.ReturnCode.Failure - end - SciMLBase.build_linear_solution(alg, y, nothing, cache; retcode = retcode) - =# - SciMLBase.build_linear_solution(alg, y, nothing, cache) - end -end - -@static if INCLUDE_SPARSE - include("factorization_sparse.jl") -end - -# Solver Specific Traits -## Needs Square Matrix -""" - needs_square_A(alg) - -Returns `true` if the algorithm requires a square matrix. -""" -needs_square_A(::Nothing) = false # Linear Solve automatically will use a correct alg! -needs_square_A(alg::SciMLLinearSolveAlgorithm) = true -for alg in (:QRFactorization, :FastQRFactorization, :NormalCholeskyFactorization, - :NormalBunchKaufmanFactorization) - @eval needs_square_A(::$(alg)) = false -end -for kralg in (Krylov.lsmr!, Krylov.craigmr!) - @eval needs_square_A(::KrylovJL{$(typeof(kralg))}) = false -end -for alg in (:LUFactorization, :FastLUFactorization, :SVDFactorization, - :GenericFactorization, :GenericLUFactorization, :SimpleLUFactorization, - :RFLUFactorization, :UMFPACKFactorization, :KLUFactorization, :SparspakFactorization, - :DiagonalFactorization, :CholeskyFactorization, :BunchKaufmanFactorization, - :CHOLMODFactorization, :LDLtFactorization, :AppleAccelerateLUFactorization, - :MKLLUFactorization, :MetalLUFactorization) - @eval needs_square_A(::$(alg)) = true -end - -const IS_OPENBLAS = Ref(true) -isopenblas() = IS_OPENBLAS[] - -const HAS_APPLE_ACCELERATE = Ref(false) -appleaccelerate_isavailable() = HAS_APPLE_ACCELERATE[] - -PrecompileTools.@compile_workload begin - A = rand(4, 4) - b = rand(4) - prob = LinearProblem(A, b) - sol = solve(prob) - sol = solve(prob, LUFactorization()) - sol = solve(prob, RFLUFactorization()) - sol = solve(prob, KrylovJL_GMRES()) -end - -@static if INCLUDE_SPARSE - PrecompileTools.@compile_workload begin - A = sprand(4, 4, 0.3) + I - b = rand(4) - prob = LinearProblem(A, b) - sol = solve(prob, KLUFactorization()) - sol = solve(prob, UMFPACKFactorization()) - end -end - -PrecompileTools.@compile_workload begin - A = sprand(4, 4, 0.3) + I - b = rand(4) - prob = LinearProblem(A * A', b) - sol = solve(prob) # in case sparspak is used as default - sol = solve(prob, SparspakFactorization()) -end - -ALREADY_WARNED_CUDSS = Ref{Bool}(false) -error_no_cudss_lu(A) = nothing -cudss_loaded(A) = false - -export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization, - GenericLUFactorization, SimpleLUFactorization, RFLUFactorization, - NormalCholeskyFactorization, NormalBunchKaufmanFactorization, - UMFPACKFactorization, KLUFactorization, FastLUFactorization, FastQRFactorization, - SparspakFactorization, DiagonalFactorization, CholeskyFactorization, - BunchKaufmanFactorization, CHOLMODFactorization, LDLtFactorization - -export LinearSolveFunction, DirectLdiv! - -export KrylovJL, KrylovJL_CG, KrylovJL_MINRES, KrylovJL_GMRES, - KrylovJL_BICGSTAB, KrylovJL_LSMR, KrylovJL_CRAIGMR, - IterativeSolversJL, IterativeSolversJL_CG, IterativeSolversJL_GMRES, - IterativeSolversJL_BICGSTAB, IterativeSolversJL_MINRES, IterativeSolversJL_IDRS, - KrylovKitJL, KrylovKitJL_CG, KrylovKitJL_GMRES - -export SimpleGMRES - -export HYPREAlgorithm -export CudaOffloadFactorization -export MKLPardisoFactorize, MKLPardisoIterate -export PardisoJL -export MKLLUFactorization -export AppleAccelerateLUFactorization -export MetalLUFactorization - -export OperatorAssumptions, OperatorCondition - -export LinearSolveAdjoint - -end +module LinearSolve +if isdefined(Base, :Experimental) && + isdefined(Base.Experimental, Symbol("@max_methods")) + @eval Base.Experimental.@max_methods 1 +end + +import PrecompileTools + +PrecompileTools.@recompile_invalidations begin + using ArrayInterface + using RecursiveFactorization + using Base: cache_dependencies, Bool + using LinearAlgebra + using SparseArrays + using SparseArrays: AbstractSparseMatrixCSC, nonzeros, rowvals, getcolptr + using LazyArrays: @~, BroadcastArray + using SciMLBase: AbstractLinearAlgorithm + using SciMLOperators + using SciMLOperators: AbstractSciMLOperator, IdentityOperator + using Setfield + using UnPack + using KLU + using Sparspak + using FastLapackInterface + using DocStringExtensions + using EnumX + using Markdown + using ChainRulesCore + import InteractiveUtils + + import StaticArraysCore: StaticArray, SVector, MVector, SMatrix, MMatrix + + using LinearAlgebra: BlasInt, LU + using LinearAlgebra.LAPACK: require_one_based_indexing, + chkfinite, chkstride1, + @blasfunc, chkargsok + + import GPUArraysCore + import Preferences + import ConcreteStructs: @concrete + + # wrap + import Krylov + using SciMLBase + import Preferences +end + +const CRC = ChainRulesCore + +if Preferences.@load_preference("LoadMKL_JLL", true) + using MKL_jll + const usemkl = MKL_jll.is_available() +else + const usemkl = false +end + +using Reexport +@reexport using SciMLBase +using SciMLBase: _unwrap_val + +abstract type SciMLLinearSolveAlgorithm <: SciMLBase.AbstractLinearAlgorithm end +abstract type AbstractFactorization <: SciMLLinearSolveAlgorithm end +abstract type AbstractKrylovSubspaceMethod <: SciMLLinearSolveAlgorithm end +abstract type AbstractSolveFunction <: SciMLLinearSolveAlgorithm end + +# Traits + +needs_concrete_A(alg::AbstractFactorization) = true +needs_concrete_A(alg::AbstractKrylovSubspaceMethod) = false +needs_concrete_A(alg::AbstractSolveFunction) = false + +# Util +is_underdetermined(x) = false +is_underdetermined(A::AbstractMatrix) = size(A, 1) < size(A, 2) +is_underdetermined(A::AbstractSciMLOperator) = size(A, 1) < size(A, 2) + +_isidentity_struct(A) = false +_isidentity_struct(λ::Number) = isone(λ) +_isidentity_struct(A::UniformScaling) = isone(A.λ) +_isidentity_struct(::SciMLOperators.IdentityOperator) = true + +# Dispatch Friendly way to check if an extension is loaded +__is_extension_loaded(::Val) = false + +# Check if a sparsity pattern has changed +pattern_changed(fact, A) = false + +function _fast_sym_givens! end + +# Code + +const INCLUDE_SPARSE = Preferences.@load_preference("include_sparse", Base.USE_GPL_LIBS) + +EnumX.@enumx DefaultAlgorithmChoice begin + LUFactorization + QRFactorization + DiagonalFactorization + DirectLdiv! + SparspakFactorization + KLUFactorization + UMFPACKFactorization + KrylovJL_GMRES + GenericLUFactorization + RFLUFactorization + LDLtFactorization + BunchKaufmanFactorization + CHOLMODFactorization + SVDFactorization + CholeskyFactorization + NormalCholeskyFactorization + AppleAccelerateLUFactorization + MKLLUFactorization + QRFactorizationPivoted + KrylovJL_CRAIGMR + KrylovJL_LSMR +end + +struct DefaultLinearSolver <: SciMLLinearSolveAlgorithm + alg::DefaultAlgorithmChoice.T +end + +const BLASELTYPES = Union{Float32, Float64, ComplexF32, ComplexF64} + +include("common.jl") +include("factorization.jl") +include("appleaccelerate.jl") +include("mkl.jl") +include("simplelu.jl") +include("simplegmres.jl") +include("iterative_wrappers.jl") +include("preconditioners.jl") +include("solve_function.jl") +include("default.jl") +include("init.jl") +include("extension_algs.jl") +include("adjoint.jl") +include("deprecated.jl") + +@inline function _issucess(F::LinearAlgebra.QRCompactWY) + (m, n) = size(F) + U = view(F.factors, 1:min(m, n), 1:n) + return any(iszero, Iterators.reverse(@view U[diagind(U)])) +end +@inline _issucess(F) = hasmethod(LinearAlgebra.issuccess, (typeof(F),)) ? + LinearAlgebra.issuccess(F) : true + +@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 _issucess(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) + end +end + +@static if INCLUDE_SPARSE + include("factorization_sparse.jl") +end + +# Solver Specific Traits +## Needs Square Matrix +""" + needs_square_A(alg) + +Returns `true` if the algorithm requires a square matrix. +""" +needs_square_A(::Nothing) = false # Linear Solve automatically will use a correct alg! +needs_square_A(alg::SciMLLinearSolveAlgorithm) = true +for alg in (:QRFactorization, :FastQRFactorization, :NormalCholeskyFactorization, + :NormalBunchKaufmanFactorization) + @eval needs_square_A(::$(alg)) = false +end +for kralg in (Krylov.lsmr!, Krylov.craigmr!) + @eval needs_square_A(::KrylovJL{$(typeof(kralg))}) = false +end +for alg in (:LUFactorization, :FastLUFactorization, :SVDFactorization, + :GenericFactorization, :GenericLUFactorization, :SimpleLUFactorization, + :RFLUFactorization, :UMFPACKFactorization, :KLUFactorization, :SparspakFactorization, + :DiagonalFactorization, :CholeskyFactorization, :BunchKaufmanFactorization, + :CHOLMODFactorization, :LDLtFactorization, :AppleAccelerateLUFactorization, + :MKLLUFactorization, :MetalLUFactorization) + @eval needs_square_A(::$(alg)) = true +end + +const IS_OPENBLAS = Ref(true) +isopenblas() = IS_OPENBLAS[] + +const HAS_APPLE_ACCELERATE = Ref(false) +appleaccelerate_isavailable() = HAS_APPLE_ACCELERATE[] + +PrecompileTools.@compile_workload begin + A = rand(4, 4) + b = rand(4) + prob = LinearProblem(A, b) + sol = solve(prob) + sol = solve(prob, LUFactorization()) + sol = solve(prob, RFLUFactorization()) + sol = solve(prob, KrylovJL_GMRES()) +end + +@static if INCLUDE_SPARSE + PrecompileTools.@compile_workload begin + A = sprand(4, 4, 0.3) + I + b = rand(4) + prob = LinearProblem(A, b) + sol = solve(prob, KLUFactorization()) + sol = solve(prob, UMFPACKFactorization()) + end +end + +PrecompileTools.@compile_workload begin + A = sprand(4, 4, 0.3) + I + b = rand(4) + prob = LinearProblem(A * A', b) + sol = solve(prob) # in case sparspak is used as default + sol = solve(prob, SparspakFactorization()) +end + +ALREADY_WARNED_CUDSS = Ref{Bool}(false) +error_no_cudss_lu(A) = nothing +cudss_loaded(A) = false + +export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization, + GenericLUFactorization, SimpleLUFactorization, RFLUFactorization, + NormalCholeskyFactorization, NormalBunchKaufmanFactorization, + UMFPACKFactorization, KLUFactorization, FastLUFactorization, FastQRFactorization, + SparspakFactorization, DiagonalFactorization, CholeskyFactorization, + BunchKaufmanFactorization, CHOLMODFactorization, LDLtFactorization + +export LinearSolveFunction, DirectLdiv! + +export KrylovJL, KrylovJL_CG, KrylovJL_MINRES, KrylovJL_GMRES, + KrylovJL_BICGSTAB, KrylovJL_LSMR, KrylovJL_CRAIGMR, + IterativeSolversJL, IterativeSolversJL_CG, IterativeSolversJL_GMRES, + IterativeSolversJL_BICGSTAB, IterativeSolversJL_MINRES, IterativeSolversJL_IDRS, + KrylovKitJL, KrylovKitJL_CG, KrylovKitJL_GMRES + +export SimpleGMRES + +export HYPREAlgorithm +export CudaOffloadFactorization +export MKLPardisoFactorize, MKLPardisoIterate +export PardisoJL +export MKLLUFactorization +export AppleAccelerateLUFactorization +export MetalLUFactorization + +export OperatorAssumptions, OperatorCondition + +export LinearSolveAdjoint + +end diff --git a/src/default.jl b/src/default.jl index 47fe280d7..4621edcff 100644 --- a/src/default.jl +++ b/src/default.jl @@ -216,9 +216,19 @@ function defaultalg(A, b, assump::OperatorAssumptions{Bool}) elseif assump.condition === OperatorCondition.WellConditioned DefaultAlgorithmChoice.NormalCholeskyFactorization elseif assump.condition === OperatorCondition.IllConditioned - DefaultAlgorithmChoice.QRFactorizationPivoted + if is_underdetermined(A) + # Underdetermined + DefaultAlgorithmChoice.QRFactorizationPivoted + else + DefaultAlgorithmChoice.QRFactorization + end elseif assump.condition === OperatorCondition.VeryIllConditioned - DefaultAlgorithmChoice.QRFactorizationPivoted + if is_underdetermined(A) + # Underdetermined + DefaultAlgorithmChoice.QRFactorizationPivoted + else + DefaultAlgorithmChoice.QRFactorization + end elseif assump.condition === OperatorCondition.SuperIllConditioned DefaultAlgorithmChoice.SVDFactorization else diff --git a/test/default_algs.jl b/test/default_algs.jl index d9fb79141..e6d69fe52 100644 --- a/test/default_algs.jl +++ b/test/default_algs.jl @@ -33,7 +33,7 @@ solve(prob) @test LinearSolve.defaultalg(nothing, zeros(5), LinearSolve.OperatorAssumptions(false)).alg === - LinearSolve.DefaultAlgorithmChoice.QRFactorizationPivoted + LinearSolve.DefaultAlgorithmChoice.QRFactorization A = spzeros(2, 2) # test that solving a singular problem doesn't error @@ -121,7 +121,7 @@ b = [1.0, 0.0, 0.0] prob = LinearProblem(A, b) sol = solve(prob) -@test sol.u ≈ [0.4, 0.2] +@test !SciMLBase.successful_retcode(sol.retcode) ## Show that we cannot select a default alg once by checking the rank, since it might change ## later in the cache @@ -143,4 +143,4 @@ cache.A = [2.0 1.0 sol = solve!(cache) -@test sol.u ≈ [0.4, 0.2] +@test !SciMLBase.successful_retcode(sol.retcode)