Merge pull request #37 from vpuri3/vp-pardiso

Pardiso

Author: ChrisRackauckas

Project.toml

@@ -11,6 +11,7 @@ KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -32,6 +33,7 @@ julia = "1.6"
 DiffEqProblemLibrary = "a077e3f3-b75c-5d7f-a0c6-6bc4c8ec64a9"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2"
 
 [targets]
-test = ["Test", "OrdinaryDiffEq", "DiffEqProblemLibrary"]
+test = ["Test", "OrdinaryDiffEq", "DiffEqProblemLibrary", "Pardiso"]

src/LinearSolve.jl

@@ -10,10 +10,11 @@ using SparseArrays
 using SciMLBase: AbstractDiffEqOperator, AbstractLinearAlgorithm
 using Setfield
 using UnPack
+using Requires
 
 # wrap
 import Krylov
-import KrylovKit
+import KrylovKit # TODO
 import IterativeSolvers
 
 using Reexport
@@ -38,12 +39,13 @@ function __init__()
   else
     IS_OPENBLAS[] = occursin("openblas", BLAS.get_config().loaded_libs[1].libname)
   end
+
+  @require Pardiso="46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" include("pardiso.jl")
 end
 
 export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization,
        RFLUFactorizaation
-export KrylovJL, KrylovJL_CG, KrylovJL_GMRES, KrylovJL_BICGSTAB,
-       KrylovJL_MINRES,
+export KrylovJL, KrylovJL_CG, KrylovJL_GMRES, KrylovJL_BICGSTAB, KrylovJL_MINRES,
        IterativeSolversJL, IterativeSolversJL_CG, IterativeSolversJL_GMRES,
        IterativeSolversJL_BICGSTAB, IterativeSolversJL_MINRES
 export DefaultLinSolve

src/pardiso.jl

@@ -0,0 +1,95 @@
+
+## Pardiso
+
+import Pardiso
+
+export PardisoJL, MKLPardisoFactorize, MKLPardisoIterate
+
+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...)
+
+# 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
+
+    if A isa DiffEqArrayOperator
+        A = A.A
+    end
+
+    solver =
+    if Pardiso.PARDISO_LOADED[]
+        solver = Pardiso.PardisoSolver()
+        solver_type !== nothing && Pardiso.set_solver!(solver, solver_type)
+
+        solver
+    else
+        solver = Pardiso.MKLPardisoSolver()
+        nprocs !== nothing && Pardiso.set_nprocs!(solver, nprocs)
+
+        solver
+    end
+
+    Pardiso.pardisoinit(solver) # default initialization
+
+    matrix_type !== nothing && Pardiso.set_matrixtype!(solver, matrix_type)
+    cache.verbose && Pardiso.set_msglvl!(solver, Pardiso.MESSAGE_LEVEL_ON)
+
+    # pass in vector of tuples like [(iparm::Int, key::Int) ...]
+    if iparm !== nothing
+        for i in length(iparm)
+            Pardiso.set_iparm!(solver, iparm[i]...)
+        end
+    end
+
+    if dparm !== nothing
+        for i in length(dparm)
+            Pardiso.set_dparm!(solver, dparm[i]...)
+        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
+
+    return solver
+end
+
+function SciMLBase.solve(cache::LinearCache, alg::PardisoJL; kwargs...)
+    @unpack A, b, u = cache
+    if A isa DiffEqArrayOperator
+        A = A.A
+    end
+
+    if cache.isfresh
+        solver = init_cacheval(alg, cache)
+        cache = set_cacheval(cache, solver)
+    end
+
+    alg.solve_phase !== nothing && Pardiso.set_phase!(cache.cacheval, alg.solve_phase)
+    Pardiso.pardiso(cache.cacheval, u, A, b)
+    alg.release_phase !== nothing && Pardiso.set_phase!(cache.cacheval, alg.release_phase)
+
+    return cache.u
+end

src/wrappers.jl

@@ -297,3 +297,4 @@ function SciMLBase.solve(cache::LinearCache, alg::IterativeSolversJL; kwargs...)
 
     return cache.u
 end
+

test/runtests.jl

@@ -126,6 +126,44 @@ end
     end
 end
 
+@testset "PardisoJL" begin
+    @test_throws UndefVarError alg = PardisoJL()
+
+    using Pardiso, SparseArrays
+
+    A1 = sparse([ 1. 0 -2  3
+                 0  5  1  2
+                -2  1  4 -7
+                 3  2 -7  5 ])
+    b1 = rand(4)
+    prob1 = LinearProblem(A1, b1)
+
+    lambda = 3
+    e = ones(n)
+    e2 = ones(n-1)
+    A2 = spdiagm(-1 => im*e2, 0 => lambda*e, 1 => -im*e2)
+    b2 = rand(n) + im * zeros(n)
+
+    prob2 = LinearProblem(A2, b2)
+
+    for alg in (
+                PardisoJL(),
+                MKLPardisoFactorize(),
+                MKLPardisoIterate(),
+               )
+
+        u = solve(prob1, alg; cache_kwargs...)
+        @test A1 * u ≈ b1
+
+        # common interface doesn't support complex types
+        # https://github.com/SciML/LinearSolve.jl/issues/38
+
+#       u = solve(prob2, alg; cache_kwargs...)
+#       @test A2 * u ≈ b2
+    end
+
+end
+
 @testset "Preconditioners" begin
     @testset "scaling_preconditioner" begin
         s = rand()