Use KrylovPreconditioners

exanauts · Nov 29, 2023 · 935de9d · 935de9d
1 parent c0d69a6
commit 935de9d
Show file tree

Hide file tree

Showing 10 changed files with 25 additions and 424 deletions.
diff --git a/Project.toml b/Project.toml
@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
+KrylovPreconditioners = "45d422c2-293f-44ce-8315-2cb988662dec"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -30,6 +31,7 @@ CUDA = "4.1, 5"
 ForwardDiff = "0.10"
 KernelAbstractions = "0.9"
 Krylov = "0.9"
+KrylovPreconditioners = "0.1"
 LazyArtifacts = "1.9"
 LightGraphs = "1.3"
 LinearAlgebra = "1.9"

diff --git a/ext/ExaPFAMDGPUExt.jl b/ext/ExaPFAMDGPUExt.jl
@@ -36,5 +36,4 @@ List all linear solvers available solving the power flow on an NVIDIA GPU.
 ExaPF.list_solvers(::ROCBackend) = [LS.BICGSTAB, LS.DQGMRES, LS.EigenBICGSTAB, LS.KrylovBICGSTAB]
 
 include("amdgpu_wrapper.jl")
-include("amdgpu_preconditioner.jl")
 end
diff --git a/ext/ExaPFCUDAExt.jl b/ext/ExaPFCUDAExt.jl
@@ -31,5 +31,4 @@ List all linear solvers available solving the power flow on an NVIDIA GPU.
 ExaPF.list_solvers(::CUDABackend) = [LS.DirectSolver, LS.BICGSTAB, LS.DQGMRES, LS.EigenBICGSTAB, LS.KrylovBICGSTAB]
 
 include("cuda_wrapper.jl")
-include("cuda_preconditioner.jl")
 end
diff --git a/ext/amdgpu_preconditioner.jl b/ext/amdgpu_preconditioner.jl
diff --git a/ext/cuda_preconditioner.jl b/ext/cuda_preconditioner.jl
diff --git a/src/LinearSolvers/LinearSolvers.jl b/src/LinearSolvers/LinearSolvers.jl
@@ -17,6 +17,7 @@ import ..ExaPF: xnorm
 
 import Base.size, Base.sizeof, Base.format_bytes
 import Krylov: KrylovStats, allocate_if, ksizeof, FloatOrComplex, ktimer, matrix_to_vector, kdisplay, mulorldiv!
+using KrylovPreconditioners
 
 const KA = KernelAbstractions
 
@@ -32,7 +33,6 @@ export DirectSolver, BICGSTAB, EigenBICGSTAB, KrylovBICGSTAB
     Diverged,
 )
 
-include("preconditioners.jl")
 include("bicgstab.jl")
 include("bicgstab_eigen.jl")
 
@@ -144,7 +144,7 @@ Custom BICGSTAB implementation to solve iteratively the linear system
 ``A  x = y``.
 """
 struct BICGSTAB <: AbstractIterativeLinearSolver
-    precond::AbstractPreconditioner
+    precond::AbstractKrylovPreconditioner
     maxiter::Int
     tol::Float64
     verbose::Bool
@@ -174,7 +174,7 @@ Julia's port of Eigen's BICGSTAB to solve iteratively the linear system
 ``A x = y``.
 """
 struct EigenBICGSTAB <: AbstractIterativeLinearSolver
-    precond::AbstractPreconditioner
+    precond::AbstractKrylovPreconditioner
     maxiter::Int
     tol::Float64
     verbose::Bool
@@ -206,7 +206,7 @@ Wrap `Krylov.jl`'s DQGMRES algorithm to solve iteratively the linear system
 """
 struct DQGMRES <: AbstractIterativeLinearSolver
     inner::Krylov.DqgmresSolver
-    precond::AbstractPreconditioner
+    precond::AbstractKrylovPreconditioner
     memory::Int
     verbose::Bool
 end
@@ -238,7 +238,7 @@ Wrap `Krylov.jl`'s BICGSTAB algorithm to solve iteratively the linear system
 """
 struct KrylovBICGSTAB <: AbstractIterativeLinearSolver
     inner::Krylov.BicgstabSolver
-    precond::AbstractPreconditioner
+    precond::AbstractKrylovPreconditioner
     verbose::Int
     atol::Float64
     rtol::Float64
@@ -255,7 +255,6 @@ end
 function ldiv!(solver::KrylovBICGSTAB,
     y::AbstractVector, J::AbstractMatrix, x::AbstractVector,
 )
-    _allowscalar(J) do
         Krylov.bicgstab!(
             solver.inner, J, x;
             N=solver.precond,
@@ -264,9 +263,8 @@ function ldiv!(solver::KrylovBICGSTAB,
             verbose=solver.verbose,
             history=true,
         )
-    end
     copyto!(y, solver.inner.x)
-    return length(solver.inner.stats.residuals)
+    return solver.inner.stats.niter
 end
 
 """