Merge pull request #64 from SciML/init

make sure to init the default algorithm and move first factorization

Author: ChrisRackauckas

src/common.jl

@@ -65,7 +65,7 @@ function set_cacheval(cache::LinearCache, alg_cache)
     return cache
 end
 
-init_cacheval(alg::Union{SciMLLinearSolveAlgorithm,Nothing}, A, b, u) = nothing
+init_cacheval(alg::SciMLLinearSolveAlgorithm, A, b, u) = nothing
 
 SciMLBase.init(prob::LinearProblem, args...; kwargs...) = SciMLBase.init(prob,nothing,args...;kwargs...)
 
@@ -91,9 +91,9 @@ function SciMLBase.init(prob::LinearProblem, alg::Union{SciMLLinearSolveAlgorith
     Pl = (Pl !== nothing) ? Pl : Identity()
     Pr = (Pr !== nothing) ? Pr : Identity()
 
-    cacheval = init_cacheval(alg, A, b, u0)
-    isfresh = cacheval === nothing
-    Tc = isfresh ? Any : typeof(cacheval)
+    cacheval = init_cacheval(alg, A, b, u0, Pl, Pr, maxiters, abstol, reltol, verbose)
+    isfresh = true
+    Tc = typeof(cacheval)
 
     A = alias_A ? A : deepcopy(A)
     b = alias_b ? b : deepcopy(b)

src/default.jl

@@ -56,3 +56,58 @@ function SciMLBase.solve(cache::LinearCache, alg::Nothing,
         SciMLBase.solve(cache, alg, args...; kwargs...)
     end
 end
+
+function init_cacheval(alg::Nothing, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    if A isa DiffEqArrayOperator
+        A = A.A
+    end
+
+    # Special case on Arrays: avoid BLAS for RecursiveFactorization.jl when
+    # it makes sense according to the benchmarks, which is dependent on
+    # whether MKL or OpenBLAS is being used
+    if A isa Matrix
+        if eltype(A) <: Union{Float32,Float64,ComplexF32,ComplexF64} &&
+                    ArrayInterface.can_setindex(b) && (size(A,1) <= 100 ||
+                                              (isopenblas() && size(A,1) <= 500)
+                                             )
+            alg = RFLUFactorization()
+            init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+        else
+            alg = LUFactorization()
+            init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+        end
+
+    # These few cases ensure the choice is optimal without the
+    # dynamic dispatching of factorize
+    elseif A isa Tridiagonal
+        alg = GenericFactorization(;fact_alg=lu!)
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    elseif A isa SymTridiagonal
+        alg = GenericFactorization(;fact_alg=ldlt!)
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    elseif A isa SparseMatrixCSC
+        alg = UMFPACKFactorization()
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+
+    # This catches the cases where a factorization overload could exist
+    # For example, BlockBandedMatrix
+    elseif ArrayInterface.isstructured(A)
+        alg = GenericFactorization()
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+
+    # This catches the case where A is a CuMatrix
+    # Which does not have LU fully defined
+    elseif !(A isa AbstractDiffEqOperator)
+        alg = QRFactorization()
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+
+    # Not factorizable operator, default to only using A*x
+    # IterativeSolvers is faster on CPU but not GPU-compatible
+    elseif cache.u isa Array
+        alg = IterativeSolversJL_GMRES()
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    else
+        alg = KrylovJL_GMRES()
+        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    end
+end

src/factorization.jl

@@ -1,13 +1,17 @@
 function SciMLBase.solve(cache::LinearCache, alg::AbstractFactorization; kwargs...)
     if cache.isfresh
-        fact = init_cacheval(alg, cache.A, cache.b, cache.u)
+        fact = do_factorization(alg, cache.A, cache.b, cache.u)
         cache = set_cacheval(cache, fact)
     end
 
     y = ldiv!(cache.u, cache.cacheval, cache.b)
     SciMLBase.build_linear_solution(alg,y,nothing,cache)
 end
 
+# Bad fallback: will fail if `A` is just a stand-in
+# This should instead just create the factorization type.
+init_cacheval(alg::AbstractFactorization, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose) = do_factorization(alg, A, b, u)
+
 ## LU Factorizations
 
 struct LUFactorization{P} <: AbstractFactorization
@@ -23,7 +27,7 @@ function LUFactorization()
     LUFactorization(pivot)
 end
 
-function init_cacheval(alg::LUFactorization, A, b, u)
+function do_factorization(alg::LUFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("LU is not defined for $(typeof(A))")
 
@@ -34,12 +38,14 @@ function init_cacheval(alg::LUFactorization, A, b, u)
     return fact
 end
 
+init_cacheval(alg::LUFactorization, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose) = ArrayInterface.lu_instance(A)
+
 # This could be a GenericFactorization perhaps?
 Base.@kwdef struct UMFPACKFactorization <: AbstractFactorization
     reuse_symbolic::Bool = true
 end
 
-function init_cacheval(::UMFPACKFactorization, A, b, u)
+function do_factorization(::UMFPACKFactorization, A, b, u)
     if A isa AbstractDiffEqOperator
         A = A.A
     end
@@ -62,7 +68,7 @@ function SciMLBase.solve(cache::LinearCache, alg::UMFPACKFactorization)
             SuiteSparse.UMFPACK.umfpack_symbolic!(cache.cacheval)
             fact = lu!(cache.cacheval, A)
         else
-            fact = init_cacheval(alg, A, cache.b, cache.u)
+            fact = do_factorization(alg, A, cache.b, cache.u)
         end
         cache = set_cacheval(cache, fact)
     end
@@ -75,7 +81,7 @@ Base.@kwdef struct KLUFactorization <: AbstractFactorization
     reuse_symbolic::Bool = true
 end
 
-function init_cacheval(::KLUFactorization, A, b, u)
+function do_factorization(::KLUFactorization, A, b, u)
     if A isa AbstractDiffEqOperator
         A = A.A
     end
@@ -98,7 +104,7 @@ function SciMLBase.solve(cache::LinearCache, alg::KLUFactorization)
             KLU.klu_analyze!(cache.cacheval)
             fact = klu!(cache.cacheval, A)
         else
-            fact = init_cacheval(alg, A, cache.b, cache.u)
+            fact = do_factorization(alg, A, cache.b, cache.u)
         end
         cache = set_cacheval(cache, fact)
     end
@@ -123,7 +129,7 @@ function QRFactorization()
     QRFactorization(pivot, 16)
 end
 
-function init_cacheval(alg::QRFactorization, A, b, u)
+function do_factorization(alg::QRFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("QR is not defined for $(typeof(A))")
 
@@ -143,7 +149,7 @@ end
 
 SVDFactorization() = SVDFactorization(false, LinearAlgebra.DivideAndConquer())
 
-function init_cacheval(alg::SVDFactorization, A, b, u)
+function do_factorization(alg::SVDFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("SVD is not defined for $(typeof(A))")
 
@@ -164,7 +170,7 @@ end
 GenericFactorization(;fact_alg = LinearAlgebra.factorize) =
     GenericFactorization(fact_alg)
 
-function init_cacheval(alg::GenericFactorization, A, b, u)
+function do_factorization(alg::GenericFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("GenericFactorization is not defined for $(typeof(A))")
 

src/iterative_wrappers.jl

@@ -69,7 +69,7 @@ function get_KrylovJL_solver(KrylovAlg)
     return KS
 end
 
-function init_cacheval(alg::KrylovJL, A, b, u)
+function init_cacheval(alg::KrylovJL, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
 
     KS = get_KrylovJL_solver(alg.KrylovAlg)
 
@@ -101,7 +101,7 @@ end
 
 function SciMLBase.solve(cache::LinearCache, alg::KrylovJL; kwargs...)
     if cache.isfresh
-        solver = init_cacheval(alg, cache.A, cache.b, cache.u)
+        solver = init_cacheval(alg, cache.A, cache.b, cache.u, cache.Pl, cache.Pr, cache.maxiters, cache.abstol, cache.reltol, cache.verbose)
         cache = set_cacheval(cache, solver)
     end
 
@@ -183,20 +183,13 @@ IterativeSolversJL_MINRES(args...;kwargs...) =
                        generate_iterator=IterativeSolvers.minres_iterable!,
                        kwargs...)
 
-function init_cacheval(alg::IterativeSolversJL, cache::LinearCache)
-    @unpack A, b, u = cache
-
-    Pl = get_preconditioner(alg.Pl, cache.Pl)
-    Pr = get_preconditioner(alg.Pr, cache.Pr)
-
-    abstol  = cache.abstol
-    reltol  = cache.reltol
-    maxiter = cache.maxiters
-    verbose = cache.verbose
+function init_cacheval(alg::IterativeSolversJL, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    Pl = get_preconditioner(alg.Pl, Pl)
+    Pr = get_preconditioner(alg.Pr, Pr)
 
     restart = (alg.gmres_restart == 0) ? min(20, size(A,1)) : alg.gmres_restart
 
-    kwargs = (abstol=abstol, reltol=reltol, maxiter=maxiter,
+    kwargs = (abstol=abstol, reltol=reltol, maxiter=maxiters,
               alg.kwargs...)
 
     iterable = if alg.generate_iterator === IterativeSolvers.cg_iterator!
@@ -212,19 +205,19 @@ function init_cacheval(alg::IterativeSolversJL, cache::LinearCache)
           @warn "$(alg.generate_iterator) doesn't support right preconditioning"
         alg.generate_iterator(u, A, b, alg.args...; Pl=Pl,
                               abstol=abstol, reltol=reltol,
-                              max_mv_products=maxiter*2,
+                              max_mv_products=maxiters*2,
                               alg.kwargs...)
     else # minres, qmr
         alg.generate_iterator(u, A, b, alg.args...;
-                              abstol=abstol, reltol=reltol, maxiter=maxiter,
+                              abstol=abstol, reltol=reltol, maxiter=maxiters,
                               alg.kwargs...)
     end
     return iterable
 end
 
 function SciMLBase.solve(cache::LinearCache, alg::IterativeSolversJL; kwargs...)
     if cache.isfresh
-        solver = init_cacheval(alg, cache)
+        solver = init_cacheval(alg, cache.A, cache.b, cache.u, cache.Pl, cache.Pr, cache.maxiters, cache.abstol, cache.reltol, cache.verbose)
         cache = set_cacheval(cache, solver)
     end
 

src/pardiso.jl

@@ -2,24 +2,17 @@ Base.@kwdef struct PardisoJL <: SciMLLinearSolveAlgorithm
     nprocs::Union{Int, Nothing} = nothing
     solver_type::Union{Int, Pardiso.Solver, Nothing} = nothing
     matrix_type::Union{Int, Pardiso.MatrixType, Nothing} = nothing
-    fact_phase::Union{Int, Pardiso.Phase, Nothing} = nothing
-    solve_phase::Union{Int, Pardiso.Phase, Nothing} = nothing
-    release_phase::Union{Int, Nothing} = nothing
     iparm::Union{Vector{Tuple{Int,Int}}, Nothing} = nothing
     dparm::Union{Vector{Tuple{Int,Int}}, Nothing} = nothing
 end
 
-MKLPardisoFactorize(;kwargs...) = PardisoJL(;fact_phase=Pardiso.NUM_FACT,
-                                             solve_phase=Pardiso.SOLVE_ITERATIVE_REFINE,
-                                             kwargs...)
-MKLPardisoIterate(;kwargs...) = PardisoJL(;solve_phase=Pardiso.NUM_FACT_SOLVE_REFINE,
-                                           kwargs...)
+MKLPardisoFactorize(;kwargs...) = PardisoJL(;kwargs...)
+MKLPardisoIterate(;kwargs...) = PardisoJL(;kwargs...)
 
 # TODO schur complement functionality
 
-function init_cacheval(alg::PardisoJL, cache::LinearCache)
-    @unpack nprocs, solver_type, matrix_type, fact_phase, solve_phase, iparm, dparm = alg
-    @unpack A, b, u = cache
+function init_cacheval(alg::PardisoJL, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
+    @unpack nprocs, solver_type, matrix_type, iparm, dparm = alg
 
     if A isa DiffEqArrayOperator
         A = A.A
@@ -51,7 +44,7 @@ function init_cacheval(alg::PardisoJL, cache::LinearCache)
             error("Number type not supported by Pardiso")
         end
     end
-    cache.verbose && Pardiso.set_msglvl!(solver, Pardiso.MESSAGE_LEVEL_ON)
+    verbose && Pardiso.set_msglvl!(solver, Pardiso.MESSAGE_LEVEL_ON)
 
     # pass in vector of tuples like [(iparm::Int, key::Int) ...]
     if iparm !== nothing
@@ -66,15 +59,8 @@ function init_cacheval(alg::PardisoJL, cache::LinearCache)
         end
     end
 
-    if (fact_phase !== nothing) | (solve_phase !== nothing)
-        Pardiso.set_phase!(solver, Pardiso.ANALYSIS)
-        Pardiso.pardiso(solver, u, A, b)
-    end
-
-    if fact_phase !== nothing
-        Pardiso.set_phase!(solver, fact_phase)
-        Pardiso.pardiso(solver, u, A, b)
-    end
+    Pardiso.set_phase!(solver, Pardiso.ANALYSIS)
+    Pardiso.pardiso(solver, u, A, b)
 
     return solver
 end
@@ -86,15 +72,17 @@ function SciMLBase.solve(cache::LinearCache, alg::PardisoJL; kwargs...)
     end
 
     if cache.isfresh
-        solver = init_cacheval(alg, cache)
-        cache = set_cacheval(cache, solver)
+        Pardiso.set_phase!(cache.cacheval, Pardiso.NUM_FACT)
+        Pardiso.pardiso(cache.cacheval, cache.u, cache.A, cache.b)
     end
 
-    alg.solve_phase !== nothing && Pardiso.set_phase!(cache.cacheval, alg.solve_phase)
+    Pardiso.set_phase!(cache.cacheval, Pardiso.SOLVE_ITERATIVE_REFINE)
     Pardiso.pardiso(cache.cacheval, u, A, b)
-    alg.release_phase !== nothing && Pardiso.set_phase!(cache.cacheval, alg.release_phase)
 
     return SciMLBase.build_linear_solution(alg,cache.u,nothing,cache)
 end
 
+# Add finalizer to release memory
+# Pardiso.set_phase!(cache.cacheval, Pardiso.RELEASE_ALL)
+
 export PardisoJL, MKLPardisoFactorize, MKLPardisoIterate

test/runtests.jl

@@ -43,24 +43,24 @@ end
     y = solve(prob1)
     @test A1 *  y  ≈ b1
 
-    _prob = LinearProblem(SymTridiagonal(A1.A), b1; u0=x1)
+    _prob = LinearProblem(SymTridiagonal(A1), b1; u0=x1)
     y = solve(_prob)
     @test A1 *  y  ≈ b1
 
-    _prob = LinearProblem(Tridiagonal(A1.A), b1; u0=x1)
+    _prob = LinearProblem(Tridiagonal(A1), b1; u0=x1)
     y = solve(_prob)
     @test A1 *  y  ≈ b1
 
-    _prob = LinearProblem(Symmetric(A1.A), b1; u0=x1)
+    _prob = LinearProblem(Symmetric(A1), b1; u0=x1)
     y = solve(_prob)
     @test A1 *  y  ≈ b1
 
-    _prob = LinearProblem(Hermitian(A1.A), b1; u0=x1)
+    _prob = LinearProblem(Hermitian(A1), b1; u0=x1)
     y = solve(_prob)
     @test A1 *  y  ≈ b1
 
 
-    _prob = LinearProblem(sparse(A1.A), b1; u0=x1)
+    _prob = LinearProblem(sparse(A1), b1; u0=x1)
     y = solve(_prob)
     @test A1 *  y  ≈ b1
 end