Skip to content

Commit 1850ac4

Browse files
author
Jeremy E Kozdon
committed
Add (basic) ksp functionality
1 parent af7ca14 commit 1850ac4

6 files changed

Lines changed: 231 additions & 76 deletions

File tree

‎src/PETSc.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ include("matshell.jl")
2929
# include("dm.jl")
3030
# include("dmda.jl")
3131
# include("dmstag.jl")
32-
# include("ksp.jl")
32+
include("ksp.jl")
3333
# include("pc.jl")
3434
# include("snes.jl")
3535
include("sys.jl")

‎src/ksp.jl‎

Lines changed: 133 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,146 @@
1-
21
const CKSP = Ptr{Cvoid}
32
const CKSPType = Cstring
43

5-
abstract type AbstractKSP{T, PetscLib} <: Factorization{T} end
4+
abstract type AbstractKSP{PetscLib, PetscScalar} <: Factorization{PetscScalar} end
65

7-
Base.@kwdef mutable struct KSP{T, PetscLib} <: AbstractKSP{T, PetscLib}
6+
Base.@kwdef mutable struct KSP{PetscLib, PetscScalar} <:
7+
AbstractKSP{PetscLib, PetscScalar}
88
ptr::CKSP = C_NULL
9-
opts::Options{PetscLib}
10-
# Stuff to keep around so that they don't get gc'ed
11-
_A = nothing
12-
_P = nothing
13-
_dm = nothing
14-
# Function pointers
15-
ComputeRHS! = nothing
16-
ComputeOperators! = nothing
9+
opts::Options{PetscLib} = Options(PetscLib)
10+
A::Union{AbstractMat, Nothing} = nothing
11+
P::Union{AbstractMat, Nothing} = nothing
12+
end
13+
14+
"""
15+
KSP(A::AbstractMat, P::AbstractMat{PetscLib} = A; options...)
16+
17+
Create a `KSP` using the matrix `A` and preconditioner construction matrix `P`
18+
with the `options`.
19+
20+
The communicator is obtained from `A` and if it has size `1` then the garbage
21+
collector is set, otherwise the user is responsible for calling
22+
[`destroy`](@ref).
23+
24+
# External Links
25+
$(_doc_external("KSP/KSPCreate"))
26+
$(_doc_external("KSP/KSPSetOperators"))
27+
$(_doc_external("KSP/KSPSetFromOptions"))
28+
"""
29+
function KSP(
30+
A::AbstractMat{PetscLib},
31+
P::AbstractMat{PetscLib} = A;
32+
options...,
33+
) where {PetscLib}
34+
@assert initialized(PetscLib)
35+
opts = Options(PetscLib; options...)
36+
PetscScalar = PetscLib.PetscScalar
37+
ksp = KSP{PetscLib, PetscScalar}(opts = opts)
38+
comm = getcomm(A)
39+
40+
with(ksp.opts) do
41+
LibPETSc.KSPCreate(PetscLib, comm, ksp)
42+
end
43+
44+
setoperators!(ksp, A, P)
45+
setfromoptions!(ksp)
46+
47+
# If there is only one rank we can finalize the KSP with GC
48+
if MPI.Comm_size(comm) == 1
49+
finalizer(destroy, ksp)
50+
end
51+
52+
return ksp
53+
end
54+
55+
function setoperators!(
56+
ksp::AbstractKSP{PetscLib},
57+
A::AbstractMat{PetscLib},
58+
P::AbstractMat{PetscLib} = A,
59+
) where {PetscLib}
60+
LibPETSc.KSPSetOperators(PetscLib, ksp, A, P)
61+
ksp.A = A
62+
ksp.P = P
63+
return ksp
64+
end
65+
66+
function setfromoptions!(ksp::AbstractKSP{PetscLib}) where {PetscLib}
67+
with(ksp.opts) do
68+
LibPETSc.KSPSetFromOptions(PetscLib, ksp)
69+
end
70+
end
71+
72+
function destroy(ksp::AbstractKSP{PetscLib}) where {PetscLib}
73+
finalized(PetscLib) || LibPETSc.MatDestroy(PetscLib, ksp)
74+
return nothing
75+
end
76+
77+
function solve!(
78+
x::AbstractVec{PetscLib},
79+
ksp::AbstractKSP{PetscLib},
80+
b::AbstractVec{PetscLib},
81+
) where {PetscLib}
82+
with(ksp.opts) do
83+
LibPETSc.KSPSolve(PetscLib, ksp, b, x)
84+
end
85+
return x
86+
end
87+
88+
"""
89+
createvecs(ksp::AbstractKSP; nright = 0, nleft = 0)
90+
91+
Create `nright` right and `nleft` left vectors compatible with the `ksp`.
92+
Returned object `V` has `Tuple` members `V.right` and `V.left` containing the
93+
vectors.
94+
95+
# External Links
96+
$(_doc_external("KSP/KSPCreateVecs"))
97+
"""
98+
function createvecs(
99+
ksp::AbstractKSP{PetscLib};
100+
nright = 0,
101+
nleft = 0,
102+
) where {PetscLib}
103+
# pointer of pointers to the base vectors
104+
r_right_vs = Ref{Ptr{CVec}}()
105+
r_left_vs = Ref{Ptr{CVec}}()
106+
107+
# create 1 right and left vector
108+
LibPETSc.KSPCreateVecs(PetscLib, ksp, 1, r_right_vs, 1, r_left_vs)
109+
110+
# create right vectors
111+
a_v = unsafe_wrap(Array, r_right_vs[], 1; own = false)
112+
v = VecPtr(PetscLib, a_v[1], false)
113+
right = ntuple(i -> similar(v), nright)
114+
115+
# create left vectors
116+
a_v = unsafe_wrap(Array, r_left_vs[], 1; own = false)
117+
v = VecPtr(PetscLib, a_v[1], false)
118+
left = ntuple(i -> similar(v), nleft)
119+
120+
LibPETSc.VecDestroyVecs(PetscLib, 1, r_right_vs)
121+
LibPETSc.VecDestroyVecs(PetscLib, 1, r_left_vs)
122+
123+
(right = right, left = left)
124+
end
125+
126+
function LinearAlgebra.ldiv!(x::AbstractVec, ksp::AbstractKSP, b::AbstractVec)
127+
solve!(x, ksp, b)
128+
end
129+
130+
function Base.:\(ksp::AbstractKSP, b::AbstractVec)
131+
x = createvecs(ksp; nleft = 1).left[1]
132+
ldiv!(x, ksp, b)
133+
return x
17134
end
18135

136+
#=
137+
#
138+
# OLD WRAPPERS
139+
#
19140
struct WrappedKSP{T, PetscLib} <: AbstractKSP{T, PetscLib}
20141
ptr::CKSP
21142
end
22143
23-
scalartype(::KSP{T}) where {T} = T
24-
Base.eltype(::KSP{T}) where {T} = T
25-
26144
LinearAlgebra.transpose(ksp) = LinearAlgebra.Transpose(ksp)
27145
LinearAlgebra.adjoint(ksp) = LinearAlgebra.Adjoint(ksp)
28146
@@ -70,32 +188,6 @@ struct Fn_KSPComputeOperators{T} end
70188
71189
@for_libpetsc begin
72190
73-
function KSP{$PetscScalar}(comm::MPI.Comm; kwargs...)
74-
@assert initialized($petsclib)
75-
opts = Options($petsclib, kwargs...)
76-
ksp = KSP{$PetscScalar, $PetscLib}(opts=opts)
77-
with(ksp.opts) do
78-
@chk ccall((:KSPCreate, $libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CKSP}), comm, ksp)
79-
end
80-
if comm == MPI.COMM_SELF
81-
finalizer(destroy, ksp)
82-
end
83-
return ksp
84-
end
85-
86-
function destroy(ksp::KSP{$PetscScalar})
87-
finalized($petsclib) ||
88-
@chk ccall((:KSPDestroy, $libpetsc), PetscErrorCode, (Ptr{CKSP},), ksp)
89-
return nothing
90-
end
91-
92-
function setoperators!(ksp::KSP{$PetscScalar}, A::AbstractMat{$PetscScalar}, P::AbstractMat{$PetscScalar})
93-
@chk ccall((:KSPSetOperators, $libpetsc), PetscErrorCode, (CKSP, CMat, CMat), ksp, A, P)
94-
ksp._A = A
95-
ksp._P = P
96-
return nothing
97-
end
98-
99191
function (::Fn_KSPComputeRHS{$PetscScalar})(
100192
new_ksp_ptr::CKSP,
101193
cb::CVec,
@@ -184,12 +276,6 @@ struct Fn_KSPComputeOperators{T} end
184276
return nothing
185277
end
186278
187-
function setfromoptions!(ksp::KSP{$PetscScalar})
188-
with(ksp.opts) do
189-
@chk ccall((:KSPSetFromOptions, $libpetsc), PetscErrorCode, (CKSP,), ksp)
190-
end
191-
end
192-
193279
function gettype(ksp::KSP{$PetscScalar})
194280
t_r = Ref{CKSPType}()
195281
@chk ccall((:KSPGetType, $libpetsc), PetscErrorCode, (CKSP, Ptr{CKSPType}), ksp, t_r)
@@ -217,14 +303,6 @@ struct Fn_KSPComputeOperators{T} end
217303
return r_rnorm[]
218304
end
219305
220-
function solve!(x::AbstractVec{$PetscScalar}, ksp::KSP{$PetscScalar}, b::AbstractVec{$PetscScalar})
221-
with(ksp.opts) do
222-
@chk ccall((:KSPSolve, $libpetsc), PetscErrorCode,
223-
(CKSP, CVec, CVec), ksp, b, x)
224-
end
225-
return x
226-
end
227-
228306
function solve!(ksp::KSP{$PetscScalar})
229307
with(ksp.opts) do
230308
@chk ccall((:KSPSolve, $libpetsc), PetscErrorCode,
@@ -256,21 +334,6 @@ function LinearAlgebra.ldiv!(x::AbstractVector{T}, ksp::KSPAT{T, LT}, b::Abstrac
256334
end
257335
Base.:\(ksp::KSPAT{T, LT}, b::AbstractVector{T}) where {T, LT} = ldiv!(similar(b), ksp, b)
258336
259-
260-
"""
261-
KSP(A, P; options...)
262-
263-
Construct a PETSc Krylov subspace solver.
264-
265-
Any PETSc options prefixed with `ksp_` and `pc_` can be passed as keywords.
266-
"""
267-
function KSP(A::AbstractMat{T}, P::AbstractMat{T}=A; kwargs...) where {T}
268-
ksp = KSP{T}(getcomm(A); kwargs...)
269-
setoperators!(ksp, A, P)
270-
setfromoptions!(ksp)
271-
return ksp
272-
end
273-
274337
"""
275338
KSP(da::AbstractDM; options...)
276339
@@ -290,7 +353,6 @@ end
290353
291354
Base.show(io::IO, ksp::KSP) = _show(io, ksp)
292355
293-
294356
"""
295357
iters(ksp::KSP)
296358
@@ -301,7 +363,6 @@ $(_doc_external("KSP/KSPGetIterationNumber"))
301363
"""
302364
iters
303365
304-
305366
"""
306367
resnorm(ksp::KSP)
307368
@@ -311,4 +372,4 @@ Gets the last (approximate preconditioned) residual norm that has been computed.
311372
$(_doc_external("KSP/KSPGetResidualNorm"))
312373
"""
313374
resnorm
314-
375+
=#

‎src/sys.jl‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
const CPetscObject = Ptr{Cvoid}
22

3-
const UnionPetscTypes = Union{Options, AbstractVec, AbstractMat}
3+
const UnionPetscTypes = Union{Options, AbstractVec, AbstractMat, AbstractKSP}
44

55
# allows us to pass PETSc_XXX objects directly into CXXX ccall signatures
66
Base.cconvert(::Type{CPetscObject}, obj::UnionPetscTypes) = obj
@@ -12,7 +12,11 @@ function Base.unsafe_convert(::Type{Ptr{CPetscObject}}, obj::UnionPetscTypes)
1212
end
1313

1414
function getcomm(
15-
obj::Union{AbstractVec{PetscLib}, AbstractMat{PetscLib}},
15+
obj::Union{
16+
AbstractVec{PetscLib},
17+
AbstractMat{PetscLib},
18+
AbstractKSP{PetscLib},
19+
},
1620
) where {PetscLib}
1721
comm = MPI.Comm()
1822
LibPETSc.PetscObjectGetComm(PetscLib, obj, comm)

‎src/vec.jl‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -612,3 +612,10 @@ function getpetsctype(vec::AbstractVec{PetscLib}) where {PetscLib}
612612
LibPETSc.VecGetType(PetscLib, vec, name_r)
613613
return unsafe_string(name_r[])
614614
end
615+
616+
function Base.similar(v::AbstractVec{PetscLib}) where {PetscLib}
617+
r_x = Ref{CVec}()
618+
LibPETSc.VecDuplicate(PetscLib, v, r_x)
619+
x = VecPtr(PetscLib, r_x[], true)
620+
return x
621+
end

‎test/ksp.jl‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
using Test
2+
using MPI
3+
MPI.Initialized() || MPI.Init()
4+
using PETSc
5+
using LinearAlgebra: mul!
6+
7+
@testset "KSP" begin
8+
comm = MPI.COMM_WORLD
9+
mpisize = MPI.Comm_size(comm)
10+
mpirank = MPI.Comm_rank(comm)
11+
12+
for petsclib in PETSc.petsclibs
13+
PETSc.initialize(petsclib)
14+
PetscScalar = petsclib.PetscScalar
15+
PetscInt = petsclib.PetscInt
16+
17+
loc_num_rows = 10
18+
loc_num_cols = 10
19+
diag_nonzeros = 3
20+
off_diag_non_zeros = 3
21+
22+
A = PETSc.MatAIJ(
23+
petsclib,
24+
comm,
25+
loc_num_rows,
26+
loc_num_cols,
27+
diag_nonzeros,
28+
off_diag_non_zeros,
29+
)
30+
31+
# Get compatible vectors
32+
(x, b) = PETSc.createvecs(A)
33+
34+
row_rng = PETSc.ownershiprange(A, false)
35+
for i in row_rng
36+
if i == 0
37+
vals = [-2, 1]
38+
row0idxs = [i]
39+
col0idxs = [i, i + 1]
40+
elseif i == mpisize * loc_num_rows - 1
41+
vals = [-2, 1]
42+
row0idxs = [i]
43+
col0idxs = [i, i - 1]
44+
else
45+
vals = [1, -2, 1]
46+
row0idxs = [i]
47+
col0idxs = [i - 1, i, i + 1]
48+
end
49+
PETSc.setvalues!(
50+
A,
51+
PetscInt.(row0idxs),
52+
PetscInt.(col0idxs),
53+
PetscScalar.(vals),
54+
)
55+
x[i + 1] = (i + 1)^3
56+
end
57+
PETSc.assemble!(A)
58+
PETSc.assemble!(x)
59+
60+
mul!(b, A, x)
61+
y = similar(x)
62+
63+
ksp = PETSc.KSP(A; ksp_rtol = 1e-16, pc_type = "jacobi")
64+
PETSc.solve!(y, ksp, b)
65+
PETSc.withlocalarray!(x, y) do x, y
66+
@test x ≈ y
67+
end
68+
PETSc.destroy(x)
69+
70+
x = ksp \ b
71+
PETSc.withlocalarray!(x, y) do x, y
72+
@test x ≈ y
73+
end
74+
PETSc.destroy(x)
75+
76+
# PETSc.destroy(ksp)
77+
PETSc.destroy(A)
78+
PETSc.destroy(y)
79+
PETSc.destroy(b)
80+
PETSc.finalize(petsclib)
81+
end
82+
end

0 commit comments

Comments
 (0)