From 0f22d51c10e658bbbb6ad626982f7d269c8b4336 Mon Sep 17 00:00:00 2001 From: Jeremy E Kozdon Date: Tue, 3 Aug 2021 12:54:49 -0700 Subject: [PATCH] WIP: Working on KSP --- src/ksp.jl | 130 ++++++++++++++++++++++------------------------- src/options.jl | 9 ++-- src/sys.jl | 5 +- src/vec.jl | 7 +++ test/ksp.jl | 76 +++++++++++++++++++++++++++ test/runtests.jl | 3 +- 6 files changed, 155 insertions(+), 75 deletions(-) create mode 100644 test/ksp.jl diff --git a/src/ksp.jl b/src/ksp.jl index e2bd2628..b521f88d 100644 --- a/src/ksp.jl +++ b/src/ksp.jl @@ -2,37 +2,88 @@ const CKSP = Ptr{Cvoid} const CKSPType = Cstring abstract type AbstractKSP{PetscLib, PetscScalar} <: Factorization{PetscScalar} end - -Base.@kwdef mutable struct KSP{PetscLib, PetscScalar} <: AbstractKSP{PetscLib, PetscScalar} + +Base.@kwdef mutable struct KSP{PetscLib, PetscScalar} <: + AbstractKSP{PetscLib, PetscScalar} ptr::CKSP = C_NULL opts::Options{PetscLib} = Options(PetscLib) A::Union{AbstractMat, Nothing} = nothing + P::Union{AbstractMat, Nothing} = nothing end -function KSP(A::AbstractMat{PetscLib}; kwargs...) where PetscLib +function KSP( + A::AbstractMat{PetscLib}, + P::AbstractMat{PetscLib} = A; + kwargs..., +) where {PetscLib} @assert initialized(PetscLib) opts = Options(PetscLib; kwargs...) PetscScalar = PetscLib.PetscScalar - ksp = KSP{PetscLib, PetscScalar}(opts=opts, A = A) - #= + ksp = KSP{PetscLib, PetscScalar}(opts = opts) + comm = getcomm(A) + with(ksp.opts) do - @chk ccall((:KSPCreate, $libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CKSP}), comm, ksp) + LibPETSc.KSPCreate(PetscLib, comm, ksp) end - if comm == MPI.COMM_SELF + + setoperators!(ksp, A, P) + setfromoptions!(ksp) + + # If there is only one rank we can finalize the KSP with GC + if MPI.Comm_size(comm) == 1 finalizer(destroy, ksp) end - =# + + return ksp +end + +function setoperators!( + ksp::AbstractKSP{PetscLib}, + A::AbstractMat{PetscLib}, + P::AbstractMat{PetscLib} = A, +) where {PetscLib} + LibPETSc.KSPSetOperators(PetscLib, ksp, A, P) + ksp.A = A + ksp.P = P return ksp end +function setfromoptions!(ksp::AbstractKSP{PetscLib}) where {PetscLib} + with(ksp.opts) do + LibPETSc.KSPSetFromOptions(PetscLib, ksp) + end +end + +function destroy(ksp::AbstractKSP{PetscLib}) where {PetscLib} + finalized(PetscLib) || LibPETSc.MatDestroy(PetscLib, ksp) + return nothing +end + +function solve!( + x::AbstractVec{PetscLib}, + ksp::AbstractKSP{PetscLib}, + b::AbstractVec{PetscLib}, +) where {PetscLib} + with(ksp.opts) do + LibPETSc.KSPSolve(PetscLib, ksp, b, x) + end + return x +end + +function LinearAlgebra.ldiv!(x::AbstractVec, ksp::AbstractKSP, b::AbstractVec) + solve!(x, ksp, b) +end +#= +function Base.:\(ksp::AbstractKSP, b::AbstractVec) + ldiv!(similar(b), ksp, b) +end +=# + #= struct WrappedKSP{T, PetscLib} <: AbstractKSP{T, PetscLib} ptr::CKSP end -scalartype(::KSP{T}) where {T} = T -Base.eltype(::KSP{T}) where {T} = T - LinearAlgebra.transpose(ksp) = LinearAlgebra.Transpose(ksp) LinearAlgebra.adjoint(ksp) = LinearAlgebra.Adjoint(ksp) @@ -80,32 +131,6 @@ struct Fn_KSPComputeOperators{T} end @for_libpetsc begin - function KSP{$PetscScalar}(comm::MPI.Comm; kwargs...) - @assert initialized($petsclib) - opts = Options($petsclib, kwargs...) - ksp = KSP{$PetscScalar, $PetscLib}(opts=opts) - with(ksp.opts) do - @chk ccall((:KSPCreate, $libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CKSP}), comm, ksp) - end - if comm == MPI.COMM_SELF - finalizer(destroy, ksp) - end - return ksp - end - - function destroy(ksp::KSP{$PetscScalar}) - finalized($petsclib) || - @chk ccall((:KSPDestroy, $libpetsc), PetscErrorCode, (Ptr{CKSP},), ksp) - return nothing - end - - function setoperators!(ksp::KSP{$PetscScalar}, A::AbstractMat{$PetscScalar}, P::AbstractMat{$PetscScalar}) - @chk ccall((:KSPSetOperators, $libpetsc), PetscErrorCode, (CKSP, CMat, CMat), ksp, A, P) - ksp._A = A - ksp._P = P - return nothing - end - function (::Fn_KSPComputeRHS{$PetscScalar})( new_ksp_ptr::CKSP, cb::CVec, @@ -194,12 +219,6 @@ struct Fn_KSPComputeOperators{T} end return nothing end - function setfromoptions!(ksp::KSP{$PetscScalar}) - with(ksp.opts) do - @chk ccall((:KSPSetFromOptions, $libpetsc), PetscErrorCode, (CKSP,), ksp) - end - end - function gettype(ksp::KSP{$PetscScalar}) t_r = Ref{CKSPType}() @chk ccall((:KSPGetType, $libpetsc), PetscErrorCode, (CKSP, Ptr{CKSPType}), ksp, t_r) @@ -227,14 +246,6 @@ struct Fn_KSPComputeOperators{T} end return r_rnorm[] end - function solve!(x::AbstractVec{$PetscScalar}, ksp::KSP{$PetscScalar}, b::AbstractVec{$PetscScalar}) - with(ksp.opts) do - @chk ccall((:KSPSolve, $libpetsc), PetscErrorCode, - (CKSP, CVec, CVec), ksp, b, x) - end - return x - end - function solve!(ksp::KSP{$PetscScalar}) with(ksp.opts) do @chk ccall((:KSPSolve, $libpetsc), PetscErrorCode, @@ -266,21 +277,6 @@ function LinearAlgebra.ldiv!(x::AbstractVector{T}, ksp::KSPAT{T, LT}, b::Abstrac end Base.:\(ksp::KSPAT{T, LT}, b::AbstractVector{T}) where {T, LT} = ldiv!(similar(b), ksp, b) - -""" - KSP(A, P; options...) - -Construct a PETSc Krylov subspace solver. - -Any PETSc options prefixed with `ksp_` and `pc_` can be passed as keywords. -""" -function KSP(A::AbstractMat{T}, P::AbstractMat{T}=A; kwargs...) where {T} - ksp = KSP{T}(getcomm(A); kwargs...) - setoperators!(ksp, A, P) - setfromoptions!(ksp) - return ksp -end - """ KSP(da::AbstractDM; options...) @@ -300,7 +296,6 @@ end Base.show(io::IO, ksp::KSP) = _show(io, ksp) - """ iters(ksp::KSP) @@ -311,7 +306,6 @@ $(_doc_external("KSP/KSPGetIterationNumber")) """ iters - """ resnorm(ksp::KSP) diff --git a/src/options.jl b/src/options.jl index 22780ba7..5652db9f 100644 --- a/src/options.jl +++ b/src/options.jl @@ -77,7 +77,7 @@ mutable struct Options{T} <: AbstractOptions{T} ptr::CPetscOptions end -function Options(petsclib::PetscLibType) +function Options_(petsclib::PetscLibType) @assert initialized(petsclib) PetscLib = typeof(petsclib) opts = Options{PetscLib}(C_NULL) @@ -86,9 +86,10 @@ function Options(petsclib::PetscLibType) return opts end -Options(petsclib; kwargs...) = Options(petsclib, kwargs...) -function Options(petsclib, ps::Pair...) - opts = Options(petsclib) +Options(petsclib::PetscLibType; kwargs...) = Options_(petsclib, kwargs...) +Options(PetscLib::Type{<:PetscLibType}; kwargs...) = Options_(getlib(PetscLib), kwargs...) +function Options_(petsclib::PetscLibType, ps::Pair...) + opts = Options_(petsclib) for (k, v) in ps opts[k] = v end diff --git a/src/sys.jl b/src/sys.jl index 88380c5f..5bfaedb6 100644 --- a/src/sys.jl +++ b/src/sys.jl @@ -1,6 +1,6 @@ const CPetscObject = Ptr{Cvoid} -const UnionPetscTypes = Union{Options, AbstractVec, AbstractMat} +const UnionPetscTypes = Union{Options, AbstractVec, AbstractMat, AbstractKSP} # allows us to pass PETSc_XXX objects directly into CXXX ccall signatures Base.cconvert(::Type{CPetscObject}, obj::UnionPetscTypes) = obj @@ -12,7 +12,8 @@ function Base.unsafe_convert(::Type{Ptr{CPetscObject}}, obj::UnionPetscTypes) end function getcomm( - obj::Union{AbstractVec{PetscLib}, AbstractMat{PetscLib}}, + obj::Union{AbstractVec{PetscLib}, AbstractMat{PetscLib}, + AbstractKSP{PetscLib}}, ) where {PetscLib} comm = MPI.Comm() LibPETSc.PetscObjectGetComm(PetscLib, obj, comm) diff --git a/src/vec.jl b/src/vec.jl index 1d69e74c..21258134 100644 --- a/src/vec.jl +++ b/src/vec.jl @@ -612,3 +612,10 @@ function getpetsctype(vec::AbstractVec{PetscLib}) where {PetscLib} LibPETSc.VecGetType(PetscLib, vec, name_r) return unsafe_string(name_r[]) end + +function Base.similar(v::AbstractVec{PetscLib}) where {PetscLib} + r_x = Ref{CVec}() + LibPETSc.VecDuplicate(PetscLib, v, r_x) + x = VecPtr(PetscLib, r_x[], true) + return x +end diff --git a/test/ksp.jl b/test/ksp.jl new file mode 100644 index 00000000..9ffec6dc --- /dev/null +++ b/test/ksp.jl @@ -0,0 +1,76 @@ +using Test +using MPI +MPI.Initialized() || MPI.Init() +using PETSc +using LinearAlgebra: mul! + +@testset "KSP" begin + comm = MPI.COMM_WORLD + mpisize = MPI.Comm_size(comm) + mpirank = MPI.Comm_rank(comm) + + for petsclib in PETSc.petsclibs + PETSc.initialize(petsclib) + PetscScalar = petsclib.PetscScalar + PetscInt = petsclib.PetscInt + + loc_num_rows = 10 + loc_num_cols = 10 + diag_nonzeros = 3 + off_diag_non_zeros = 3 + + A = PETSc.MatAIJ( + petsclib, + comm, + loc_num_rows, + loc_num_cols, + diag_nonzeros, + off_diag_non_zeros, + ) + + # Get compatible vectors + (x, b) = PETSc.createvecs(A) + + row_rng = PETSc.ownershiprange(A, false) + for i in row_rng + if i == 0 + vals = [-2, 1] + row0idxs = [i] + col0idxs = [i, i + 1] + elseif i == mpisize * loc_num_rows - 1 + vals = [-2, 1] + row0idxs = [i] + col0idxs = [i, i - 1] + else + vals = [1, -2, 1] + row0idxs = [i] + col0idxs = [i - 1, i, i + 1] + end + PETSc.setvalues!( + A, + PetscInt.(row0idxs), + PetscInt.(col0idxs), + PetscScalar.(vals), + ) + x[i + 1] = (i + 1)^3 + end + PETSc.assemble!(A) + PETSc.assemble!(x) + + mul!(b, A, x) + y = similar(x) + + ksp = PETSc.KSP(A; ksp_rtol = 1e-16, pc_type = "jacobi") + PETSc.solve!(y, ksp, b) + PETSc.withlocalarray!(x, y) do x, y + @test x ≈ y + end + + # PETSc.destroy(ksp) + PETSc.destroy(A) + PETSc.destroy(x) + PETSc.destroy(y) + PETSc.destroy(b) + PETSc.finalize(petsclib) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 85d81e9d..9c4a43ab 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,7 +3,7 @@ using MPI: mpiexec # Do the MPI tests first so we do not have mpi running inside MPI @testset "mpi tests" begin - for file in ("mpivec.jl", "mpimat.jl") + for file in ("mpivec.jl", "mpimat.jl", "ksp.jl") @test mpiexec() do mpi_cmd cmd = `$mpi_cmd -n 4 $(Base.julia_cmd()) --startup-file=no --project $file` @@ -18,3 +18,4 @@ include("options.jl") include("vec.jl") include("mat.jl") include("matshell.jl") +include("ksp.jl")