1 |
C $Header: /u/gcmpack/MITgcm_contrib/dgoldberg/code_cg3d_petsc/cg3d_petsc.F,v 1.2 2016/07/02 17:48:26 dgoldberg Exp $ |
2 |
C $Name: $ |
3 |
|
4 |
#include "CPP_OPTIONS.h" |
5 |
|
6 |
C---+----1----+----2----+----3----+----4----+----5----+----6----+----7-|--+----| |
7 |
|
8 |
CBOP |
9 |
SUBROUTINE CG3D_PETSC( |
10 |
U cg3d_x, ! solution vector |
11 |
U cg3d_b, ! rhs |
12 |
I tolerance, |
13 |
U maxIter, |
14 |
I myIter, |
15 |
I myThid ) |
16 |
|
17 |
C /============================================================\ |
18 |
C | SUBROUTINE | |
19 |
C | o | |
20 |
C |============================================================| |
21 |
C | | |
22 |
C \============================================================/ |
23 |
IMPLICIT NONE |
24 |
|
25 |
#include "SIZE.h" |
26 |
#include "EEPARAMS.h" |
27 |
#include "PARAMS.h" |
28 |
#include "CG3D.h" |
29 |
|
30 |
#ifdef ALLOW_PETSC |
31 |
#include "finclude/petsc.h" |
32 |
#include "CG3D_PETSC.h" |
33 |
! UNCOMMENT IF V3.0 |
34 |
!#include "finclude/petscvec.h" |
35 |
!#include "finclude/petscmat.h" |
36 |
!#include "finclude/petscksp.h" |
37 |
!#include "finclude/petscpc.h" |
38 |
#endif |
39 |
C === Global variables === |
40 |
|
41 |
|
42 |
C !INPUT/OUTPUT ARGUMENTS |
43 |
C cg_Uin, cg_Vin - input and output velocities |
44 |
C cg_Bu, cg_Bv - driving stress |
45 |
INTEGER myThid |
46 |
INTEGER maxiter |
47 |
INTEGER myIter |
48 |
_RL tolerance |
49 |
_RL cg3d_x (1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr,nSx,nSy) |
50 |
_RL cg3d_b (1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr,nSx,nSy) |
51 |
|
52 |
|
53 |
C LOCAL VARIABLES |
54 |
INTEGER i, j, bi, bj, cg_halo, conv_flag |
55 |
INTEGER iter, col_iter, k |
56 |
_RL dot_p1, dot_p2, alpha_k, beta_k, resid, resid_0 |
57 |
_RL dot_p1_tile (nSx,nSy) |
58 |
_RL dot_p2_tile (nSx,nSy) |
59 |
CHARACTER*(MAX_LEN_MBUF) msgBuf |
60 |
! LOGICAL create_mat, destroy_mat |
61 |
|
62 |
|
63 |
#ifdef ALLOW_PETSC |
64 |
INTEGER indices(Nr*(snx*nsx*sny*nsy)) |
65 |
INTEGER cg3d_dofs_cum_sum (0:nPx*nPy*MAX_CG3D_PETSC_CPUINVERT-1), |
66 |
& idx(1) |
67 |
_RL rhs_values(Nr*(snx*nsx*sny*nsy)) |
68 |
_RL solution_values(Nr*(snx*nsx*sny*nsy)) |
69 |
_RL mat_values (7,1), mat_val_return(1) |
70 |
INTEGER indices_col(7) |
71 |
INTEGER local_dofs, global_dofs, dof_index, dof_index_col |
72 |
INTEGER local_offset |
73 |
INTEGER iters |
74 |
INTEGER color, rank |
75 |
PC subpc |
76 |
KSP subksp(1) |
77 |
Vec rhs |
78 |
Vec solution |
79 |
PetscErrorCode ierr |
80 |
#ifdef ALLOW_USE_MPI |
81 |
integer mpiRC, mpiMyWid |
82 |
#endif |
83 |
#endif |
84 |
|
85 |
|
86 |
! print *, "GOT HERE CG3D", myByLo(mythid), myByHi(myThid), |
87 |
! & myBxLo(mythid), myBxHi(myThid) |
88 |
! RETURN |
89 |
|
90 |
#ifdef ALLOW_PETSC |
91 |
|
92 |
IF ((myIter.eq.1) .OR. (.not.cg3d_petsc_reuse_mat)) THEN |
93 |
cg3d_need2create_mat = .TRUE. |
94 |
ELSE |
95 |
cg3d_need2create_mat = .FALSE. |
96 |
ENDIF |
97 |
|
98 |
IF ((myIter.eq.nTimeSteps) .OR. (.not.cg3d_petsc_reuse_mat)) THEN |
99 |
cg3d_need2destroy_mat = .TRUE. |
100 |
ELSE |
101 |
cg3d_need2destroy_mat = .FALSE. |
102 |
ENDIF |
103 |
|
104 |
! Print *, "GOT HERE CG3D", local_dofs,global_dofs, |
105 |
! & cg3d_dofs_cum_sum |
106 |
|
107 |
#ifdef ALLOW_USE_MPI |
108 |
|
109 |
|
110 |
CALL MPI_COMM_RANK( MPI_COMM_WORLD, mpiMyWId, mpiRC ) |
111 |
local_dofs = cg3d_dofs_process (mpiMyWid) |
112 |
global_dofs = 0 |
113 |
|
114 |
|
115 |
cg3d_dofs_cum_sum(0) = 0 |
116 |
|
117 |
DO i=0,nPx*nPy*MAX_CG3D_PETSC_CPUINVERT-1 |
118 |
! IF (i.le.cg3d_petsc_cpuInVert) THEN |
119 |
global_dofs = global_dofs + cg3d_dofs_process (i) |
120 |
if (i.ge.1) then |
121 |
cg3d_dofs_cum_sum(i) = cg3d_dofs_cum_sum(i-1)+ |
122 |
& cg3d_dofs_process(i-1) |
123 |
endif |
124 |
! ENDIF |
125 |
ENDDO |
126 |
|
127 |
local_offset = cg3d_dofs_cum_sum(mpimywid) |
128 |
|
129 |
#else |
130 |
|
131 |
local_dofs = cg3d_dofs_process (0) |
132 |
global_dofs = local_dofs |
133 |
local_offset = 0 |
134 |
|
135 |
#endif |
136 |
Print *, "GOT HERE CG3D", local_dofs,global_dofs, |
137 |
& cg3d_dofs_cum_sum |
138 |
|
139 |
!---------------------- |
140 |
|
141 |
|
142 |
call VecCreate(PETSC_COMM_WORLD, rhs, ierr) |
143 |
call VecSetSizes(rhs, local_dofs, global_dofs, ierr) |
144 |
call VecSetType(rhs, VECMPI, ierr) |
145 |
|
146 |
call VecCreate(PETSC_COMM_WORLD, solution, ierr) |
147 |
call VecSetSizes(solution, local_dofs, global_dofs, ierr) |
148 |
call VecSetType(solution, VECMPI, ierr) |
149 |
|
150 |
do i=1,local_dofs |
151 |
indices(i) = i-1 + local_offset |
152 |
end do |
153 |
do i=1,Nr*nSx*nSy*sNx*sNy |
154 |
rhs_values (i) = 0. _d 0 |
155 |
solution_values (i) = 0. _d 0 |
156 |
enddo |
157 |
|
158 |
! gather rhs and initial guess values to populate petsc vectors |
159 |
|
160 |
DO bj = myByLo(myThid), myByHi(myThid) |
161 |
DO bi = myBxLo(myThid), myBxHi(myThid) |
162 |
DO j=1,sNy |
163 |
DO i=1,sNx |
164 |
DO k=1,Nr |
165 |
|
166 |
color = INT(cg3d_petsc_color(i,j,k,bi,bj)) |
167 |
rank = cg3d_color_rank (color) |
168 |
dof_index = INT(cg3d_petsc_dofs(i,j,k,bi,bj)) |
169 |
& - local_offset |
170 |
|
171 |
if (dof_index.ge.0 .and. rank.eq.mpimywid) THEN |
172 |
|
173 |
rhs_values(dof_index+1) = cg3d_b(i,j,k,bi,bj) |
174 |
solution_values(dof_index+1) = cg3d_x(i,j,k,bi,bj) |
175 |
|
176 |
endif |
177 |
|
178 |
ENDDO |
179 |
ENDDO |
180 |
ENDDO |
181 |
ENDDO |
182 |
ENDDO |
183 |
|
184 |
call VecSetValues(rhs, local_dofs, indices, rhs_values, |
185 |
& INSERT_VALUES, ierr) |
186 |
call VecAssemblyBegin(rhs, ierr) |
187 |
call VecAssemblyEnd(rhs, ierr) |
188 |
|
189 |
|
190 |
call VecSetValues(solution, local_dofs, indices, |
191 |
& solution_values, INSERT_VALUES, ierr) |
192 |
call VecAssemblyBegin(solution, ierr) |
193 |
call VecAssemblyEnd(solution, ierr) |
194 |
|
195 |
|
196 |
|
197 |
if (cg3d_need2create_mat) then |
198 |
CALL TIMER_START ('CG3D_PETSC_MATCREATE',myThid) |
199 |
|
200 |
! IF USING v3.0 THEN |
201 |
! call MatCreateMPIAIJ (PETSC_COMM_WORLD, |
202 |
call MatCreateAIJ (PETSC_COMM_WORLD, |
203 |
& local_dofs, local_dofs, |
204 |
& global_dofs, global_dofs, |
205 |
& 7, PETSC_NULL_INTEGER, |
206 |
& 7, PETSC_NULL_INTEGER, |
207 |
& matrix_petsc_cg3d, ierr) |
208 |
|
209 |
|
210 |
! populate petsc matrix |
211 |
|
212 |
DO bj = myByLo(myThid), myByHi(myThid) |
213 |
DO bi = myBxLo(myThid), myBxHi(myThid) |
214 |
DO j=1,sNy |
215 |
DO i=1,sNx |
216 |
DO k=1,Nr |
217 |
|
218 |
dof_index = INT(cg3d_petsc_dofs(i,j,k,bi,bj)) |
219 |
! & - local_offset |
220 |
color = INT(cg3d_petsc_color(i,j,k,bi,bj)) |
221 |
rank = cg3d_color_rank (color) |
222 |
|
223 |
IF (dof_index .ge. 0 .and. rank.eq.mpimywid) THEN |
224 |
|
225 |
DO col_iter=1,7 |
226 |
indices_col(col_iter) = 0 |
227 |
mat_values(col_iter,1) = 0. _d 0 |
228 |
ENDDO |
229 |
col_iter = 0 |
230 |
|
231 |
! ASSUME THE STENCIL |
232 |
! ac3d (i,j,k) --> 1 |
233 |
! aw3d (i,j,k) --> 2 |
234 |
! as3d (i,j,k) --> 3 |
235 |
! av3d (i,j,k) --> 4 |
236 |
! aw3d (i+1,j,k) --> 5 |
237 |
! as3d (i,j+1,k) --> 6 |
238 |
! av3d (i,j,k+1) --> 7 |
239 |
|
240 |
dof_index_col = cg3d_petsc_dofs(i,j,k,bi,bj) |
241 |
if (dof_index_col.ge.0) then |
242 |
col_iter = col_iter + 1 |
243 |
mat_values(col_iter,1) = ac3d(i,j,k,bi,bj) |
244 |
indices_col(col_iter) = dof_index_col |
245 |
endif |
246 |
|
247 |
dof_index_col = cg3d_petsc_dofs(i-1,j,k,bi,bj) |
248 |
if (dof_index_col.ge.0) then |
249 |
col_iter = col_iter + 1 |
250 |
mat_values(col_iter,1) = aw3d(i,j,k,bi,bj) |
251 |
indices_col(col_iter) = dof_index_col |
252 |
endif |
253 |
|
254 |
dof_index_col = cg3d_petsc_dofs(i,j-1,k,bi,bj) |
255 |
if (dof_index_col.ge.0) then |
256 |
col_iter = col_iter + 1 |
257 |
mat_values(col_iter,1) = as3d(i,j,k,bi,bj) |
258 |
indices_col(col_iter) = dof_index_col |
259 |
endif |
260 |
|
261 |
if (k.gt.1) then |
262 |
dof_index_col = cg3d_petsc_dofs(i,j,k-1,bi,bj) |
263 |
if (dof_index_col.ge.0) then |
264 |
col_iter = col_iter + 1 |
265 |
mat_values(col_iter,1) = av3d(i,j,k,bi,bj) |
266 |
indices_col(col_iter) = dof_index_col |
267 |
endif |
268 |
endif |
269 |
|
270 |
dof_index_col = cg3d_petsc_dofs(i+1,j,k,bi,bj) |
271 |
if (dof_index_col.ge.0) then |
272 |
col_iter = col_iter + 1 |
273 |
mat_values(col_iter,1) = aw3d(i+1,j,k,bi,bj) |
274 |
indices_col(col_iter) = dof_index_col |
275 |
endif |
276 |
|
277 |
dof_index_col = cg3d_petsc_dofs(i,j+1,k,bi,bj) |
278 |
if (dof_index_col.ge.0) then |
279 |
col_iter = col_iter + 1 |
280 |
mat_values(col_iter,1) = as3d(i,j+1,k,bi,bj) |
281 |
indices_col(col_iter) = dof_index_col |
282 |
endif |
283 |
|
284 |
if (k.lt.Nr) then |
285 |
dof_index_col = cg3d_petsc_dofs(i,j,k+1,bi,bj) |
286 |
if (dof_index_col.ge.0) then |
287 |
col_iter = col_iter + 1 |
288 |
mat_values(col_iter,1) = av3d(i,j,k+1,bi,bj) |
289 |
indices_col(col_iter) = dof_index_col |
290 |
endif |
291 |
endif |
292 |
|
293 |
|
294 |
call matSetValues (matrix_petsc_cg3d, 1, dof_index, col_iter, |
295 |
& indices_col, |
296 |
& mat_values,INSERT_VALUES,ierr) |
297 |
|
298 |
|
299 |
ENDIF |
300 |
|
301 |
|
302 |
ENDDO |
303 |
ENDDO |
304 |
ENDDO |
305 |
ENDDO |
306 |
ENDDO |
307 |
|
308 |
call MatAssemblyBegin(matrix_petsc_cg3d,MAT_FINAL_ASSEMBLY,ierr) |
309 |
call MatAssemblyEnd(matrix_petsc_cg3d,MAT_FINAL_ASSEMBLY,ierr) |
310 |
|
311 |
|
312 |
call KSPCreate(PETSC_COMM_WORLD, ksp_cg3d, ierr) |
313 |
call KSPSetOperators(ksp_cg3d, |
314 |
& matrix_petsc_cg3d, matrix_petsc_cg3d, |
315 |
& DIFFERENT_NONZERO_PATTERN, ierr) |
316 |
|
317 |
IF (CG3D_PETSC_PRECOND_TYPE.eq.'MUMPS') then |
318 |
call KSPSetType(ksp_cg3d,KSPPREONLY,ierr) |
319 |
ELSE |
320 |
SELECT CASE (CG3D_PETSC_SOLVER_TYPE) |
321 |
CASE ('CG') |
322 |
PRINT *, "PETSC SOLVER: SELECTED CG" |
323 |
call KSPSetType(ksp_cg3d, KSPCG, ierr) |
324 |
CASE ('GMRES') |
325 |
PRINT *, "PETSC SOLVER: SELECTED GMRES" |
326 |
call KSPSetType(ksp_cg3d, KSPGMRES, ierr) |
327 |
CASE ('BICG') |
328 |
PRINT *, "PETSC SOLVER: SELECTED BICG" |
329 |
call KSPSetType(ksp_cg3d, KSPBICG, ierr) |
330 |
CASE DEFAULT |
331 |
PRINT *, "PETSC SOLVER: SELECTED DEFAULT" |
332 |
call KSPSetType(ksp_cg3d, KSPCG, ierr) |
333 |
END SELECT |
334 |
ENDIF |
335 |
|
336 |
call KSPGetPC(ksp_cg3d, pc_cg3d, ierr) |
337 |
call KSPSetTolerances(ksp_cg3d,tolerance, |
338 |
& PETSC_DEFAULT_DOUBLE_PRECISION, |
339 |
& PETSC_DEFAULT_DOUBLE_PRECISION, |
340 |
& maxiter,ierr) |
341 |
|
342 |
SELECT CASE (CG3D_PETSC_PRECOND_TYPE) |
343 |
CASE ('BLOCKJACOBI') |
344 |
PRINT *, "PETSC PRECOND: SELECTED BJACOBI" |
345 |
call PCSetType(pc_cg3d, PCBJACOBI, ierr) |
346 |
call kspsetup (ksp_cg3d, ierr) |
347 |
call PCBJacobiGetSubKSP (pc_cg3d,PETSC_NULL_INTEGER, |
348 |
& PETSC_NULL_INTEGER, |
349 |
& subksp,ierr); |
350 |
call KSPGetPC (subksp, subpc, ierr) |
351 |
call PCSetType (subpc, PCILU, ierr) |
352 |
! call PCFactorSetLevels(subpc,0,ierr) |
353 |
CASE ('JACOBI') |
354 |
PRINT *, "PETSC PRECOND: SELECTED JACOBI" |
355 |
call PCSetType(pc_cg3d, PCJACOBI, ierr) |
356 |
CASE ('ILU') |
357 |
PRINT *, "PETSC PRECOND: SELECTED ILU" |
358 |
call PCSetType(pc_cg3d, PCILU, ierr) |
359 |
|
360 |
CASE ('GAMG') |
361 |
PRINT *, "PETSC PRECOND: SELECTED GAMG" |
362 |
call PCSetType(pc_cg3d, PCGAMG, ierr) |
363 |
call PCGAMGSetCoarseEqLim(pc_cg3d,10000,ierr) |
364 |
call PCGAMGSetSymGraph(pc_cg3d, PETSC_TRUE,ierr) |
365 |
call PCGAMGSetNSmooths(pc_cg3d, 0,ierr) |
366 |
call PCGAMGSetThreshold(pc_cg3d, .001,ierr) |
367 |
! call PCGAMGSetReuseProl(pc,PETSC_FALSE,ierr) |
368 |
call kspsetup (ksp_cg3d, ierr) |
369 |
|
370 |
CASE ('MUMPS') |
371 |
PRINT *, "PETSC PRECOND: SELECTED MUMPS" |
372 |
call PCSetType(pc_cg3d,PCLU,ierr) |
373 |
call PCFactorSetMatSolverPackage(pc_cg3d,MATSOLVERMUMPS,ierr) |
374 |
call PCFactorSetUpMatSolverPackage(pc_cg3d,ierr) |
375 |
call PCFactorGetMatrix(pc_cg3d,mumpsFac_cg3d,ierr) |
376 |
call MatMumpsSetIcntl(mumpsfac_cg3d,7,2,ierr) |
377 |
call MatMumpsSetIcntl(mumpsfac_cg3d,24,1,ierr) |
378 |
! call MatMumpsSetCntl(mumpsfac,3,2.e-6,ierr) |
379 |
call kspsetup (ksp_cg3d, ierr) |
380 |
|
381 |
CASE DEFAULT |
382 |
PRINT *, "PETSC PRECOND: SELECTED DEFAULT" |
383 |
call PCSetType(pc_cg3d, PCBJACOBI, ierr) |
384 |
END SELECT |
385 |
|
386 |
|
387 |
CALL TIMER_STOP ('CG3D_PETSC_MATCREATE',myThid) |
388 |
endif ! need3create_mat |
389 |
|
390 |
|
391 |
CALL TIMER_START ('CG3D_PETSC_SOLVE',myThid) |
392 |
|
393 |
call KSPSolve(ksp_cg3d, rhs, solution, ierr) |
394 |
|
395 |
|
396 |
CALL TIMER_STOP ('CG3D_PETSC_SOLVE',myThid) |
397 |
|
398 |
call KSPGetIterationNumber(ksp_cg3d,iters,ierr) |
399 |
|
400 |
maxIter = iters |
401 |
|
402 |
call VecGetValues(solution,local_dofs,indices, |
403 |
& solution_values,ierr) |
404 |
|
405 |
DO bj = myByLo(myThid), myByHi(myThid) |
406 |
DO bi = myBxLo(myThid), myBxHi(myThid) |
407 |
DO j=1,sNy |
408 |
DO i=1,sNx |
409 |
DO k=1,Nr |
410 |
|
411 |
dof_index = INT(cg3d_petsc_dofs(i,j,k,bi,bj)) |
412 |
& - local_offset |
413 |
color = INT(cg3d_petsc_color(i,j,k,bi,bj)) |
414 |
rank = cg3d_color_rank (color) |
415 |
if (dof_index.ge.0.and.rank.eq.mpimywid) THEN |
416 |
cg3d_x(i,j,k,bi,bj) = solution_values(dof_index+1) |
417 |
endif |
418 |
|
419 |
ENDDO |
420 |
ENDDO |
421 |
ENDDO |
422 |
ENDDO |
423 |
ENDDO |
424 |
|
425 |
if (cg3d_need2destroy_mat) then |
426 |
call KSPDestroy (ksp_cg3d, ierr) |
427 |
call MatDestroy (matrix_petsc_cg3d, ierr) |
428 |
endif |
429 |
call VecDestroy (rhs, ierr) |
430 |
call VecDestroy (solution, ierr) |
431 |
|
432 |
|
433 |
|
434 |
|
435 |
#endif |
436 |
|
437 |
|
438 |
|
439 |
RETURN |
440 |
END |