Actual source code: bjkokkos.kokkos.cxx
1: #define PETSC_SKIP_CXX_COMPLEX_FIX // Kokkos::complex does not need the PetscComplex fix
3: #include <petsc/private/pcbjkokkosimpl.h>
4: #include <petsc/private/kspimpl.h>
5: #include <petscksp.h>
6: #include <../src/mat/impls/aij/mpi/mpiaij.h>
7: #include <../src/mat/impls/aij/seq/kokkos/aijkok.hpp>
8: #include <petscsection.h>
9: #include <petscdmcomposite.h>
10: #include <../src/mat/impls/aij/seq/aij.h>
11: #include <../src/mat/impls/aij/seq/kokkos/aijkok.hpp>
12: #include <petscdevice_cupm.h>
14: static PetscErrorCode PCBJKOKKOSCreateKSP_BJKOKKOS(PC pc)
15: {
16: const char *prefix;
17: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
18: DM dm;
20: PetscFunctionBegin;
21: PetscCall(KSPCreate(PetscObjectComm((PetscObject)pc), &jac->ksp));
22: PetscCall(KSPSetNestLevel(jac->ksp, pc->kspnestlevel));
23: PetscCall(KSPSetErrorIfNotConverged(jac->ksp, pc->erroriffailure));
24: PetscCall(PetscObjectIncrementTabLevel((PetscObject)jac->ksp, (PetscObject)pc, 1));
25: PetscCall(PCGetOptionsPrefix(pc, &prefix));
26: PetscCall(KSPSetOptionsPrefix(jac->ksp, prefix));
27: PetscCall(KSPAppendOptionsPrefix(jac->ksp, "pc_bjkokkos_"));
28: PetscCall(PCGetDM(pc, &dm));
29: if (dm) {
30: PetscCall(KSPSetDM(jac->ksp, dm));
31: PetscCall(KSPSetDMActive(jac->ksp, KSP_DMACTIVE_ALL, PETSC_FALSE));
32: }
33: jac->reason = PETSC_FALSE;
34: jac->monitor = PETSC_FALSE;
35: jac->batch_target = 0;
36: jac->rank_target = 0;
37: jac->nsolves_team = 1;
38: jac->ksp->max_it = 50; // this is really for GMRES w/o restarts
39: PetscFunctionReturn(PETSC_SUCCESS);
40: }
42: // y <-- Ax
43: KOKKOS_INLINE_FUNCTION PetscErrorCode MatMult(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, const PetscInt start, const PetscInt end, const PetscScalar *x_loc, PetscScalar *y_loc)
44: {
45: Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) {
46: int rowa = ic[rowb];
47: int n = glb_Aai[rowa + 1] - glb_Aai[rowa];
48: const PetscInt *aj = glb_Aaj + glb_Aai[rowa]; // global
49: const PetscScalar *aa = glb_Aaa + glb_Aai[rowa];
50: PetscScalar sum;
51: Kokkos::parallel_reduce(Kokkos::ThreadVectorRange(team, n), [=](const int i, PetscScalar &lsum) { lsum += aa[i] * x_loc[r[aj[i]] - start]; }, sum);
52: Kokkos::single(Kokkos::PerThread(team), [=]() { y_loc[rowb - start] = sum; });
53: });
54: team.team_barrier();
55: return PETSC_SUCCESS;
56: }
58: // temp buffer per thread with reduction at end?
59: KOKKOS_INLINE_FUNCTION PetscErrorCode MatMultTranspose(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, const PetscInt start, const PetscInt end, const PetscScalar *x_loc, PetscScalar *y_loc)
60: {
61: Kokkos::parallel_for(Kokkos::TeamVectorRange(team, end - start), [=](int i) { y_loc[i] = 0; });
62: team.team_barrier();
63: Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) {
64: int rowa = ic[rowb];
65: int n = glb_Aai[rowa + 1] - glb_Aai[rowa];
66: const PetscInt *aj = glb_Aaj + glb_Aai[rowa]; // global
67: const PetscScalar *aa = glb_Aaa + glb_Aai[rowa];
68: const PetscScalar xx = x_loc[rowb - start]; // rowb = ic[rowa] = ic[r[rowb]]
69: Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, n), [=](const int &i) {
70: PetscScalar val = aa[i] * xx;
71: Kokkos::atomic_fetch_add(&y_loc[r[aj[i]] - start], val);
72: });
73: });
74: team.team_barrier();
75: return PETSC_SUCCESS;
76: }
78: typedef struct Batch_MetaData_TAG {
79: PetscInt flops;
80: PetscInt its;
81: KSPConvergedReason reason;
82: } Batch_MetaData;
84: // Solve A(BB^-1)x = y with TFQMR. Right preconditioned to get un-preconditioned residual
85: static KOKKOS_INLINE_FUNCTION PetscErrorCode BJSolve_TFQMR(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, PetscScalar *work_space_global, const int stride_global, const int nShareVec, PetscScalar *work_space_shared, const int stride_shared, PetscReal rtol, PetscReal atol, PetscReal dtol, PetscInt maxit, Batch_MetaData *metad, const PetscInt start, const PetscInt end, const PetscScalar glb_idiag[], const PetscScalar *glb_b, PetscScalar *glb_x, bool monitor)
86: {
87: using Kokkos::parallel_for;
88: using Kokkos::parallel_reduce;
89: int Nblk = end - start, it, m, stride = stride_shared, idx = 0;
90: PetscReal dp, dpold, w, dpest, tau, psi, cm, r0;
91: const PetscScalar *Diag = &glb_idiag[start];
92: PetscScalar *ptr = work_space_shared, rho = 0.0, rhoold = 0.0, a = 0.0, s = 0.0, b, eta, etaold, psiold, cf, dpi = 0.0;
94: if (idx++ == nShareVec) {
95: ptr = work_space_global;
96: stride = stride_global;
97: }
98: PetscScalar *XX = ptr;
99: ptr += stride;
100: if (idx++ == nShareVec) {
101: ptr = work_space_global;
102: stride = stride_global;
103: }
104: PetscScalar *R = ptr;
105: ptr += stride;
106: if (idx++ == nShareVec) {
107: ptr = work_space_global;
108: stride = stride_global;
109: }
110: PetscScalar *RP = ptr;
111: ptr += stride;
112: if (idx++ == nShareVec) {
113: ptr = work_space_global;
114: stride = stride_global;
115: }
116: PetscScalar *V = ptr;
117: ptr += stride;
118: if (idx++ == nShareVec) {
119: ptr = work_space_global;
120: stride = stride_global;
121: }
122: PetscScalar *T = ptr;
123: ptr += stride;
124: if (idx++ == nShareVec) {
125: ptr = work_space_global;
126: stride = stride_global;
127: }
128: PetscScalar *Q = ptr;
129: ptr += stride;
130: if (idx++ == nShareVec) {
131: ptr = work_space_global;
132: stride = stride_global;
133: }
134: PetscScalar *P = ptr;
135: ptr += stride;
136: if (idx++ == nShareVec) {
137: ptr = work_space_global;
138: stride = stride_global;
139: }
140: PetscScalar *U = ptr;
141: ptr += stride;
142: if (idx++ == nShareVec) {
143: ptr = work_space_global;
144: stride = stride_global;
145: }
146: PetscScalar *D = ptr;
147: ptr += stride;
148: if (idx++ == nShareVec) {
149: ptr = work_space_global;
150: stride = stride_global;
151: }
152: PetscScalar *AUQ = V;
154: // init: get b, zero x
155: parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
156: int rowa = ic[rowb];
157: R[rowb - start] = glb_b[rowa];
158: XX[rowb - start] = 0;
159: });
160: team.team_barrier();
161: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += R[idx] * PetscConj(R[idx]); }, dpi);
162: team.team_barrier();
163: r0 = dp = PetscSqrtReal(PetscRealPart(dpi));
164: // diagnostics
165: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
166: if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", 0, (double)dp); });
167: #endif
168: if (dp < atol) {
169: metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
170: it = 0;
171: goto done;
172: }
173: if (0 == maxit) {
174: metad->reason = KSP_CONVERGED_ITS;
175: it = 0;
176: goto done;
177: }
179: /* Make the initial Rp = R */
180: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { RP[idx] = R[idx]; });
181: team.team_barrier();
182: /* Set the initial conditions */
183: etaold = 0.0;
184: psiold = 0.0;
185: tau = dp;
186: dpold = dp;
188: /* rhoold = (r,rp) */
189: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += R[idx] * PetscConj(RP[idx]); }, rhoold);
190: team.team_barrier();
191: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
192: U[idx] = R[idx];
193: P[idx] = R[idx];
194: T[idx] = Diag[idx] * P[idx];
195: D[idx] = 0;
196: });
197: team.team_barrier();
198: static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, V));
200: it = 0;
201: do {
202: /* s <- (v,rp) */
203: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += V[idx] * PetscConj(RP[idx]); }, s);
204: team.team_barrier();
205: if (s == 0) {
206: metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
207: goto done;
208: }
209: a = rhoold / s; /* a <- rho / s */
210: /* q <- u - a v VecWAXPY(w,alpha,x,y): w = alpha x + y. */
211: /* t <- u + q */
212: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
213: Q[idx] = U[idx] - a * V[idx];
214: T[idx] = U[idx] + Q[idx];
215: });
216: team.team_barrier();
217: // KSP_PCApplyBAorAB
218: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { T[idx] = Diag[idx] * T[idx]; });
219: team.team_barrier();
220: static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, AUQ));
221: /* r <- r - a K (u + q) */
222: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { R[idx] = R[idx] - a * AUQ[idx]; });
223: team.team_barrier();
224: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += R[idx] * PetscConj(R[idx]); }, dpi);
225: team.team_barrier();
226: dp = PetscSqrtReal(PetscRealPart(dpi));
227: for (m = 0; m < 2; m++) {
228: if (!m) w = PetscSqrtReal(dp * dpold);
229: else w = dp;
230: psi = w / tau;
231: cm = 1.0 / PetscSqrtReal(1.0 + psi * psi);
232: tau = tau * psi * cm;
233: eta = cm * cm * a;
234: cf = psiold * psiold * etaold / a;
235: if (!m) {
236: /* D = U + cf D */
237: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { D[idx] = U[idx] + cf * D[idx]; });
238: } else {
239: /* D = Q + cf D */
240: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { D[idx] = Q[idx] + cf * D[idx]; });
241: }
242: team.team_barrier();
243: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { XX[idx] = XX[idx] + eta * D[idx]; });
244: team.team_barrier();
245: dpest = PetscSqrtReal(2 * it + m + 2.0) * tau;
246: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
247: if (monitor && m == 1) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", it + 1, (double)dpest); });
248: #endif
249: if (dpest < atol) {
250: metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
251: goto done;
252: }
253: if (dpest / r0 < rtol) {
254: metad->reason = KSP_CONVERGED_RTOL_NORMAL_EQUATIONS;
255: goto done;
256: }
257: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
258: if (dpest / r0 > dtol) {
259: metad->reason = KSP_DIVERGED_DTOL;
260: Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: %d it, res=%e, r_0=%e\n", team.league_rank(), it, dpest, r0); });
261: goto done;
262: }
263: #else
264: if (dpest / r0 > dtol) {
265: metad->reason = KSP_DIVERGED_DTOL;
266: goto done;
267: }
268: #endif
269: if (it + 1 == maxit) {
270: metad->reason = KSP_CONVERGED_ITS;
271: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
272: Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: TFQMR %d:%d it, res=%e, r_0=%e r_res=%e\n", team.league_rank(), it, m, dpest, r0, dpest / r0); });
273: #endif
274: goto done;
275: }
276: etaold = eta;
277: psiold = psi;
278: }
280: /* rho <- (r,rp) */
281: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += R[idx] * PetscConj(RP[idx]); }, rho);
282: team.team_barrier();
283: if (rho == 0) {
284: metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
285: goto done;
286: }
287: b = rho / rhoold; /* b <- rho / rhoold */
288: /* u <- r + b q */
289: /* p <- u + b(q + b p) */
290: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
291: U[idx] = R[idx] + b * Q[idx];
292: Q[idx] = Q[idx] + b * P[idx];
293: P[idx] = U[idx] + b * Q[idx];
294: });
295: /* v <- K p */
296: team.team_barrier();
297: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { T[idx] = Diag[idx] * P[idx]; });
298: team.team_barrier();
299: static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, V));
301: rhoold = rho;
302: dpold = dp;
304: it++;
305: } while (it < maxit);
306: done:
307: // KSPUnwindPreconditioner
308: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { XX[idx] = Diag[idx] * XX[idx]; });
309: team.team_barrier();
310: // put x into Plex order
311: parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
312: int rowa = ic[rowb];
313: glb_x[rowa] = XX[rowb - start];
314: });
315: metad->its = it;
316: if (1) {
317: int nnz = 0;
318: parallel_reduce(Kokkos::TeamVectorRange(team, start, end), [=](const int idx, int &lsum) { lsum += (glb_Aai[idx + 1] - glb_Aai[idx]); }, nnz);
319: metad->flops = 2 * (metad->its * (10 * Nblk + 2 * nnz) + 5 * Nblk);
320: } else {
321: metad->flops = 2 * (metad->its * (10 * Nblk + 2 * 50 * Nblk) + 5 * Nblk); // guess
322: }
323: return PETSC_SUCCESS;
324: }
326: // Solve Ax = y with biCG
327: static KOKKOS_INLINE_FUNCTION PetscErrorCode BJSolve_BICG(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, PetscScalar *work_space_global, const int stride_global, const int nShareVec, PetscScalar *work_space_shared, const int stride_shared, PetscReal rtol, PetscReal atol, PetscReal dtol, PetscInt maxit, Batch_MetaData *metad, const PetscInt start, const PetscInt end, const PetscScalar glb_idiag[], const PetscScalar *glb_b, PetscScalar *glb_x, bool monitor)
328: {
329: using Kokkos::parallel_for;
330: using Kokkos::parallel_reduce;
331: int Nblk = end - start, it, stride = stride_shared, idx = 0; // start in shared mem
332: PetscReal dp, r0;
333: const PetscScalar *Di = &glb_idiag[start];
334: PetscScalar *ptr = work_space_shared, dpi = 0.0, a = 1.0, beta = 0.0, betaold = 1.0, t1, t2;
336: if (idx++ == nShareVec) {
337: ptr = work_space_global;
338: stride = stride_global;
339: }
340: PetscScalar *XX = ptr;
341: ptr += stride;
342: if (idx++ == nShareVec) {
343: ptr = work_space_global;
344: stride = stride_global;
345: }
346: PetscScalar *Rl = ptr;
347: ptr += stride;
348: if (idx++ == nShareVec) {
349: ptr = work_space_global;
350: stride = stride_global;
351: }
352: PetscScalar *Zl = ptr;
353: ptr += stride;
354: if (idx++ == nShareVec) {
355: ptr = work_space_global;
356: stride = stride_global;
357: }
358: PetscScalar *Pl = ptr;
359: ptr += stride;
360: if (idx++ == nShareVec) {
361: ptr = work_space_global;
362: stride = stride_global;
363: }
364: PetscScalar *Rr = ptr;
365: ptr += stride;
366: if (idx++ == nShareVec) {
367: ptr = work_space_global;
368: stride = stride_global;
369: }
370: PetscScalar *Zr = ptr;
371: ptr += stride;
372: if (idx++ == nShareVec) {
373: ptr = work_space_global;
374: stride = stride_global;
375: }
376: PetscScalar *Pr = ptr;
377: ptr += stride;
379: /* r <- b (x is 0) */
380: parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
381: int rowa = ic[rowb];
382: Rl[rowb - start] = Rr[rowb - start] = glb_b[rowa];
383: XX[rowb - start] = 0;
384: });
385: team.team_barrier();
386: /* z <- Br */
387: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
388: Zr[idx] = Di[idx] * Rr[idx];
389: Zl[idx] = Di[idx] * Rl[idx];
390: });
391: team.team_barrier();
392: /* dp <- r'*r */
393: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Rr[idx] * PetscConj(Rr[idx]); }, dpi);
394: team.team_barrier();
395: r0 = dp = PetscSqrtReal(PetscRealPart(dpi));
396: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
397: if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", 0, (double)dp); });
398: #endif
399: if (dp < atol) {
400: metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
401: it = 0;
402: goto done;
403: }
404: if (0 == maxit) {
405: metad->reason = KSP_CONVERGED_ITS;
406: it = 0;
407: goto done;
408: }
410: it = 0;
411: do {
412: /* beta <- r'z */
413: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += Zr[idx] * PetscConj(Rl[idx]); }, beta);
414: team.team_barrier();
415: #if PCBJKOKKOS_VERBOSE_LEVEL >= 6
416: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
417: Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%7d beta = Z.R = %22.14e \n", i, (double)beta); });
418: #endif
419: #endif
420: if (beta == 0.0) {
421: metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
422: goto done;
423: }
424: if (it == 0) {
425: /* p <- z */
426: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
427: Pr[idx] = Zr[idx];
428: Pl[idx] = Zl[idx];
429: });
430: } else {
431: t1 = beta / betaold;
432: /* p <- z + b* p */
433: t2 = PetscConj(t1);
434: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
435: Pr[idx] = t1 * Pr[idx] + Zr[idx];
436: Pl[idx] = t2 * Pl[idx] + Zl[idx];
437: });
438: }
439: team.team_barrier();
440: betaold = beta;
441: /* z <- Kp */
442: static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, Pr, Zr));
443: static_cast<void>(MatMultTranspose(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, Pl, Zl));
444: /* dpi <- z'p */
445: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Zr[idx] * PetscConj(Pl[idx]); }, dpi);
446: team.team_barrier();
447: if (dpi == 0) {
448: metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
449: goto done;
450: }
451: //
452: a = beta / dpi; /* a = beta/p'z */
453: t1 = -a;
454: t2 = PetscConj(t1);
455: /* x <- x + ap */
456: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
457: XX[idx] = XX[idx] + a * Pr[idx];
458: Rr[idx] = Rr[idx] + t1 * Zr[idx];
459: Rl[idx] = Rl[idx] + t2 * Zl[idx];
460: });
461: team.team_barrier();
462: team.team_barrier();
463: /* dp <- r'*r */
464: parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Rr[idx] * PetscConj(Rr[idx]); }, dpi);
465: team.team_barrier();
466: dp = PetscSqrtReal(PetscRealPart(dpi));
467: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
468: if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", it + 1, (double)dp); });
469: #endif
470: if (dp < atol) {
471: metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
472: goto done;
473: }
474: if (dp / r0 < rtol) {
475: metad->reason = KSP_CONVERGED_RTOL_NORMAL_EQUATIONS;
476: goto done;
477: }
478: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
479: if (dp / r0 > dtol) {
480: metad->reason = KSP_DIVERGED_DTOL;
481: Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: %d it, res=%e, r_0=%e (BICG does this)\n", team.league_rank(), it, dp, r0); });
482: goto done;
483: }
484: #else
485: if (dp / r0 > dtol) {
486: metad->reason = KSP_DIVERGED_DTOL;
487: goto done;
488: }
489: #endif
490: if (it + 1 == maxit) {
491: metad->reason = KSP_CONVERGED_ITS; // don't worry about hitting max iterations
492: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
493: Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: BICG %d it, res=%e, r_0=%e r_res=%e\n", team.league_rank(), it, dp, r0, dp / r0); });
494: #endif
495: goto done;
496: }
497: /* z <- Br */
498: parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
499: Zr[idx] = Di[idx] * Rr[idx];
500: Zl[idx] = Di[idx] * Rl[idx];
501: });
503: it++;
504: } while (it < maxit);
505: done:
506: // put x back into Plex order
507: parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
508: int rowa = ic[rowb];
509: glb_x[rowa] = XX[rowb - start];
510: });
511: metad->its = it;
512: if (1) {
513: int nnz = 0;
514: parallel_reduce(Kokkos::TeamVectorRange(team, start, end), [=](const int idx, int &lsum) { lsum += (glb_Aai[idx + 1] - glb_Aai[idx]); }, nnz);
515: metad->flops = 2 * (metad->its * (10 * Nblk + 2 * nnz) + 5 * Nblk);
516: } else {
517: metad->flops = 2 * (metad->its * (10 * Nblk + 2 * 50 * Nblk) + 5 * Nblk); // guess
518: }
519: return PETSC_SUCCESS;
520: }
522: // KSP solver solve Ax = b; xout is output, bin is input
523: static PetscErrorCode PCApply_BJKOKKOS(PC pc, Vec bin, Vec xout)
524: {
525: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
526: Mat A = pc->pmat, Aseq = A;
527: PetscMPIInt rank;
529: PetscFunctionBegin;
530: PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
531: if (!A->spptr) Aseq = ((Mat_MPIAIJ *)A->data)->A; // MPI
532: PetscCall(MatSeqAIJKokkosSyncDevice(Aseq));
533: {
534: PetscInt maxit = jac->ksp->max_it;
535: const PetscInt conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0 && PCBJKOKKOS_VEC_SIZE != 1) ? PCBJKOKKOS_TEAM_SIZE : 1;
536: const PetscInt nwork = jac->nwork, nBlk = jac->nBlocks;
537: PetscScalar *glb_xdata = NULL, *dummy;
538: PetscReal rtol = jac->ksp->rtol, atol = jac->ksp->abstol, dtol = jac->ksp->divtol;
539: const PetscScalar *glb_idiag = jac->d_idiag_k->data(), *glb_bdata = NULL;
540: const PetscInt *glb_Aai, *glb_Aaj, *d_bid_eqOffset = jac->d_bid_eqOffset_k->data();
541: const PetscScalar *glb_Aaa;
542: const PetscInt *d_isicol = jac->d_isicol_k->data(), *d_isrow = jac->d_isrow_k->data();
543: PCFailedReason pcreason;
544: KSPIndex ksp_type_idx = jac->ksp_type_idx;
545: PetscMemType mtype;
546: PetscContainer container;
547: PetscInt batch_sz; // the number of repeated DMs, [DM_e_1, DM_e_2, DM_e_batch_sz, DM_i_1, ...]
548: VecScatter plex_batch = NULL; // not used
549: Vec bvec; // a copy of b for scatter (just alias to bin now)
550: PetscBool monitor = jac->monitor; // captured
551: PetscInt view_bid = jac->batch_target;
552: MatInfo info;
554: PetscCall(MatSeqAIJGetCSRAndMemType(Aseq, &glb_Aai, &glb_Aaj, &dummy, &mtype));
555: jac->max_nits = 0;
556: glb_Aaa = dummy;
557: if (jac->rank_target != rank) view_bid = -1; // turn off all but one process
558: PetscCall(MatGetInfo(A, MAT_LOCAL, &info));
559: // get field major is to map plex IO to/from block/field major
560: PetscCall(PetscObjectQuery((PetscObject)A, "plex_batch_is", (PetscObject *)&container));
561: if (container) {
562: PetscCall(VecDuplicate(bin, &bvec));
563: PetscCall(PetscContainerGetPointer(container, &plex_batch));
564: PetscCall(VecScatterBegin(plex_batch, bin, bvec, INSERT_VALUES, SCATTER_FORWARD));
565: PetscCall(VecScatterEnd(plex_batch, bin, bvec, INSERT_VALUES, SCATTER_FORWARD));
566: SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "No plex_batch_is -- require NO field major ordering for now");
567: } else {
568: bvec = bin;
569: }
570: // get x
571: PetscCall(VecGetArrayAndMemType(xout, &glb_xdata, &mtype));
572: #if PetscDefined(HAVE_CUDA)
573: PetscCheck(PetscMemTypeDevice(mtype), PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "No GPU data for x %d != %d", static_cast<int>(mtype), static_cast<int>(PETSC_MEMTYPE_DEVICE));
574: #endif
575: PetscCall(VecGetArrayReadAndMemType(bvec, &glb_bdata, &mtype));
576: #if PetscDefined(HAVE_CUDA)
577: PetscCheck(PetscMemTypeDevice(mtype), PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "No GPU data for b");
578: #endif
579: // get batch size
580: PetscCall(PetscObjectQuery((PetscObject)A, "batch size", (PetscObject *)&container));
581: if (container) {
582: PetscInt *pNf = NULL;
583: PetscCall(PetscContainerGetPointer(container, &pNf));
584: batch_sz = *pNf; // number of times to repeat the DMs
585: } else batch_sz = 1;
586: PetscCheck(nBlk % batch_sz == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "batch_sz = %" PetscInt_FMT ", nBlk = %" PetscInt_FMT, batch_sz, nBlk);
587: if (ksp_type_idx == BATCH_KSP_GMRESKK_IDX) {
588: // KK solver - move PETSc data into Kokkos Views, setup solver, solve, move data out of Kokkos, process metadata (convergence tests, etc.)
589: #if PetscDefined(HAVE_KOKKOS_KERNELS_BATCH)
590: PetscCall(PCApply_BJKOKKOSKERNELS(pc, glb_bdata, glb_xdata, glb_Aai, glb_Aaj, glb_Aaa, team_size, info, batch_sz, &pcreason));
591: #else
592: PetscCheck(ksp_type_idx != BATCH_KSP_GMRESKK_IDX, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Type: BATCH_KSP_GMRES not supported for complex");
593: #endif
594: } else { // Kokkos Krylov
595: using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space;
596: using vect2D_scr_t = Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, scr_mem_t>;
597: Kokkos::View<Batch_MetaData *, Kokkos::DefaultExecutionSpace> d_metadata("solver meta data", nBlk);
598: int stride_shared, stride_global, global_buff_words;
599: d_bid_eqOffset = jac->d_bid_eqOffset_k->data();
600: // solve each block independently
601: int scr_bytes_team_shared = 0, nShareVec = 0, nGlobBVec = 0;
602: if (jac->const_block_size) { // use shared memory for work vectors only if constant block size - TODO: test efficiency loss
603: size_t maximum_shared_mem_size = 64000;
604: PetscDevice device;
605: PetscCall(PetscDeviceGetDefault_Internal(&device));
606: PetscCall(PetscDeviceGetAttribute(device, PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK, &maximum_shared_mem_size));
607: stride_shared = jac->const_block_size; // captured
608: nShareVec = maximum_shared_mem_size / (jac->const_block_size * sizeof(PetscScalar)); // integer floor, number of vectors that fit in shared
609: if (nShareVec > nwork) nShareVec = nwork;
610: else nGlobBVec = nwork - nShareVec;
611: global_buff_words = jac->n * nGlobBVec;
612: scr_bytes_team_shared = jac->const_block_size * nShareVec * sizeof(PetscScalar);
613: } else {
614: scr_bytes_team_shared = 0;
615: stride_shared = 0;
616: global_buff_words = jac->n * nwork;
617: nGlobBVec = nwork; // not needed == fix
618: }
619: stride_global = jac->n; // captured
620: #if PetscDefined(HAVE_CUDA)
621: nvtxRangePushA("batch-kokkos-solve");
622: #endif
623: Kokkos::View<PetscScalar *, Kokkos::DefaultExecutionSpace> d_work_vecs_k("workvectors", global_buff_words); // global work vectors
624: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
625: PetscCall(PetscInfo(pc, "\tn = %d. %d shared bytes/team, %d global mem bytes, rtol=%e, num blocks %d, team_size=%d, %d vector threads, %d shared vectors, %d global vectors\n", (int)jac->n, scr_bytes_team_shared, global_buff_words, rtol, (int)nBlk, (int)team_size, PCBJKOKKOS_VEC_SIZE, nShareVec, nGlobBVec));
626: #endif
627: PetscScalar *d_work_vecs = d_work_vecs_k.data();
628: Kokkos::parallel_for(
629: "Solve", Kokkos::TeamPolicy<Kokkos::LaunchBounds<256, 4>>(nBlk, team_size, PCBJKOKKOS_VEC_SIZE).set_scratch_size(PCBJKOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes_team_shared)), KOKKOS_LAMBDA(const team_member team) {
630: const int blkID = team.league_rank(), start = d_bid_eqOffset[blkID], end = d_bid_eqOffset[blkID + 1];
631: vect2D_scr_t work_vecs_shared(team.team_scratch(PCBJKOKKOS_SHARED_LEVEL), end - start, nShareVec);
632: PetscScalar *work_buff_shared = work_vecs_shared.data();
633: PetscScalar *work_buff_global = &d_work_vecs[start]; // start inc'ed in
634: bool print = monitor && (blkID == view_bid);
635: switch (ksp_type_idx) {
636: case BATCH_KSP_BICG_IDX:
637: static_cast<void>(BJSolve_BICG(team, glb_Aai, glb_Aaj, glb_Aaa, d_isrow, d_isicol, work_buff_global, stride_global, nShareVec, work_buff_shared, stride_shared, rtol, atol, dtol, maxit, &d_metadata[blkID], start, end, glb_idiag, glb_bdata, glb_xdata, print));
638: break;
639: case BATCH_KSP_TFQMR_IDX:
640: static_cast<void>(BJSolve_TFQMR(team, glb_Aai, glb_Aaj, glb_Aaa, d_isrow, d_isicol, work_buff_global, stride_global, nShareVec, work_buff_shared, stride_shared, rtol, atol, dtol, maxit, &d_metadata[blkID], start, end, glb_idiag, glb_bdata, glb_xdata, print));
641: break;
642: default:
643: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
644: printf("Unknown KSP type %d\n", ksp_type_idx);
645: #else
646: /* void */;
647: #endif
648: }
649: });
650: Kokkos::fence();
651: #if PetscDefined(HAVE_CUDA)
652: nvtxRangePop();
653: nvtxRangePushA("Post-solve-metadata");
654: #endif
655: auto h_metadata = Kokkos::create_mirror(Kokkos::HostSpace::memory_space(), d_metadata);
656: Kokkos::deep_copy(h_metadata, d_metadata);
657: PetscInt max_nnit = -1;
658: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
659: PetscInt mbid = 0;
660: #endif
661: int out[2];
662: if (jac->reason) { // -pc_bjkokkos_ksp_converged_reason
663: #if PCBJKOKKOS_VERBOSE_LEVEL >= 3
664: #if PCBJKOKKOS_VERBOSE_LEVEL >= 4
665: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Iterations\n"));
666: #endif
667: // assume species major
668: #if PCBJKOKKOS_VERBOSE_LEVEL == 3
669: if (batch_sz != 1) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%s: max iterations per species:", ksp_type_idx == BATCH_KSP_BICG_IDX ? "bicg" : "tfqmr"));
670: else PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), " Linear solve converged due to %s iterations ", ksp_type_idx == BATCH_KSP_BICG_IDX ? "bicg" : "tfqmr"));
671: #endif
672: for (PetscInt dmIdx = 0, head = 0, s = 0; dmIdx < jac->num_dms; dmIdx += batch_sz) {
673: for (PetscInt f = 0, idx = head; f < jac->dm_Nf[dmIdx]; f++, idx++, s++) {
674: for (int bid = 0; bid < batch_sz; bid++) {
675: #if PCBJKOKKOS_VERBOSE_LEVEL >= 4
676: jac->max_nits += h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its; // report total number of iterations with high verbose
677: if (h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its > max_nnit) {
678: max_nnit = h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its;
679: mbid = bid;
680: }
681: #else
682: if (h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its > max_nnit) {
683: jac->max_nits = max_nnit = h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its;
684: mbid = bid;
685: }
686: #endif
687: }
688: #if PCBJKOKKOS_VERBOSE_LEVEL >= 4
689: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%2" PetscInt_FMT ":", s));
690: for (int bid = 0; bid < batch_sz; bid++) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3" PetscInt_FMT " ", h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its));
691: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n"));
692: #else // == 3
693: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3" PetscInt_FMT " ", max_nnit));
694: #endif
695: }
696: head += batch_sz * jac->dm_Nf[dmIdx];
697: }
698: #if PCBJKOKKOS_VERBOSE_LEVEL == 3
699: PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n"));
700: #endif
701: #endif
702: if (max_nnit == -1) { // < 3
703: for (int blkID = 0; blkID < nBlk; blkID++) {
704: if (h_metadata[blkID].its > max_nnit) {
705: jac->max_nits = max_nnit = h_metadata[blkID].its;
706: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
707: mbid = blkID;
708: #endif
709: }
710: }
711: }
712: out[0] = max_nnit;
713: out[1] = rank;
714: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, out, 1, MPI_2INT, MPI_MAXLOC, PetscObjectComm((PetscObject)A)));
715: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
716: if (0 == rank) {
717: if (batch_sz != 1)
718: PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] Linear solve converged due to %s iterations %d (max), on block %" PetscInt_FMT ", species %" PetscInt_FMT " (max)\n", out[1], KSPConvergedReasons[h_metadata[mbid].reason], out[0], mbid % batch_sz, mbid / batch_sz));
719: else PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] Linear solve converged due to %s iterations %d (max), on block %" PetscInt_FMT "\n", out[1], KSPConvergedReasons[h_metadata[mbid].reason], out[0], mbid));
720: }
721: #endif
722: }
723: for (int blkID = 0; blkID < nBlk; blkID++) {
724: PetscCall(PetscLogGpuFlops((PetscLogDouble)h_metadata[blkID].flops));
725: PetscCheck(h_metadata[blkID].reason >= 0 || !jac->ksp->errorifnotconverged, PetscObjectComm((PetscObject)pc), PETSC_ERR_CONV_FAILED, "ERROR reason=%s, its=%" PetscInt_FMT ". species %" PetscInt_FMT ", batch %" PetscInt_FMT,
726: KSPConvergedReasons[h_metadata[blkID].reason], h_metadata[blkID].its, blkID / batch_sz, blkID % batch_sz);
727: }
728: {
729: int errsum;
730: Kokkos::parallel_reduce(
731: nBlk,
732: KOKKOS_LAMBDA(const int idx, int &lsum) {
733: if (d_metadata[idx].reason < 0) ++lsum;
734: },
735: errsum);
736: pcreason = errsum ? PC_SUBPC_ERROR : PC_NOERROR;
737: if (!errsum && !jac->max_nits) { // set max its to give back to top KSP
738: for (int blkID = 0; blkID < nBlk; blkID++) {
739: if (h_metadata[blkID].its > jac->max_nits) jac->max_nits = h_metadata[blkID].its;
740: }
741: } else if (errsum) {
742: PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] ERROR Kokkos batch solver did not converge in all solves\n", (int)rank));
743: }
744: }
745: #if PetscDefined(HAVE_CUDA)
746: nvtxRangePop();
747: #endif
748: } // end of Kokkos (not Kernels) solvers block
749: PetscCall(VecRestoreArrayAndMemType(xout, &glb_xdata));
750: PetscCall(VecRestoreArrayReadAndMemType(bvec, &glb_bdata));
751: PetscCall(PCSetFailedReason(pc, pcreason));
752: // map back to Plex space - not used
753: if (plex_batch) {
754: PetscCall(VecCopy(xout, bvec));
755: PetscCall(VecScatterBegin(plex_batch, bvec, xout, INSERT_VALUES, SCATTER_REVERSE));
756: PetscCall(VecScatterEnd(plex_batch, bvec, xout, INSERT_VALUES, SCATTER_REVERSE));
757: PetscCall(VecDestroy(&bvec));
758: }
759: }
760: PetscFunctionReturn(PETSC_SUCCESS);
761: }
763: static PetscErrorCode PCSetUp_BJKOKKOS(PC pc)
764: {
765: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
766: Mat A = pc->pmat, Aseq = A; // use filtered block matrix, really "P"
767: PetscBool flg;
769: PetscFunctionBegin;
770: PetscCheck(A, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No matrix - A is used above");
771: PetscCall(PetscObjectTypeCompareAny((PetscObject)A, &flg, MATSEQAIJKOKKOS, MATMPIAIJKOKKOS, MATAIJKOKKOS, ""));
772: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "must use '-[dm_]mat_type aijkokkos -[dm_]vec_type kokkos' for -pc_type bjkokkos");
773: if (!A->spptr) Aseq = ((Mat_MPIAIJ *)A->data)->A; // MPI
774: PetscCall(MatSeqAIJKokkosSyncDevice(Aseq));
775: {
776: PetscInt Istart, Iend;
777: PetscMPIInt rank;
778: PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
779: PetscCall(MatGetOwnershipRange(A, &Istart, &Iend));
780: if (!jac->vec_diag) {
781: Vec *subX = NULL;
782: DM pack, *subDM = NULL;
783: PetscInt nDMs, n, *block_sizes = NULL;
784: IS isrow, isicol;
785: { // Permute the matrix to get a block diagonal system: d_isrow_k, d_isicol_k
786: MatOrderingType rtype;
787: const PetscInt *rowindices, *icolindices;
788: rtype = MATORDERINGRCM;
789: // get permutation. And invert. should we convert to local indices?
790: PetscCall(MatGetOrdering(Aseq, rtype, &isrow, &isicol)); // only seems to work for seq matrix
791: PetscCall(ISDestroy(&isrow));
792: PetscCall(ISInvertPermutation(isicol, PETSC_DECIDE, &isrow)); // THIS IS BACKWARD -- isrow is inverse
793: // if (rank==1) PetscCall(ISView(isicol, PETSC_VIEWER_STDOUT_SELF));
794: if (0) {
795: Mat mat_block_order; // debug
796: PetscCall(ISShift(isicol, Istart, isicol));
797: PetscCall(MatCreateSubMatrix(A, isicol, isicol, MAT_INITIAL_MATRIX, &mat_block_order));
798: PetscCall(ISShift(isicol, -Istart, isicol));
799: PetscCall(MatViewFromOptions(mat_block_order, NULL, "-ksp_batch_reorder_view"));
800: PetscCall(MatDestroy(&mat_block_order));
801: }
802: PetscCall(ISGetIndices(isrow, &rowindices)); // local idx
803: PetscCall(ISGetIndices(isicol, &icolindices));
804: const Kokkos::View<PetscInt *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_isrow_k((PetscInt *)rowindices, A->rmap->n);
805: const Kokkos::View<PetscInt *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_isicol_k((PetscInt *)icolindices, A->rmap->n);
806: jac->d_isrow_k = new Kokkos::View<PetscInt *>(Kokkos::create_mirror(DefaultMemorySpace(), h_isrow_k));
807: jac->d_isicol_k = new Kokkos::View<PetscInt *>(Kokkos::create_mirror(DefaultMemorySpace(), h_isicol_k));
808: Kokkos::deep_copy(*jac->d_isrow_k, h_isrow_k);
809: Kokkos::deep_copy(*jac->d_isicol_k, h_isicol_k);
810: PetscCall(ISRestoreIndices(isrow, &rowindices));
811: PetscCall(ISRestoreIndices(isicol, &icolindices));
812: // if (rank==1) PetscCall(ISView(isicol, PETSC_VIEWER_STDOUT_SELF));
813: }
814: // get block sizes & allocate vec_diag
815: PetscCall(PCGetDM(pc, &pack));
816: if (pack) {
817: PetscCall(PetscObjectTypeCompare((PetscObject)pack, DMCOMPOSITE, &flg));
818: if (flg) {
819: PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
820: PetscCall(DMCreateGlobalVector(pack, &jac->vec_diag));
821: } else pack = NULL; // flag for no DM
822: }
823: if (!jac->vec_diag) { // get 'nDMs' and sizes 'block_sizes' w/o DMComposite. TODO: User could provide ISs
824: PetscInt bsrt, bend, ncols, ntot = 0;
825: const PetscInt *colsA, nloc = Iend - Istart;
826: const PetscInt *rowindices, *icolindices;
827: PetscCall(PetscMalloc1(nloc, &block_sizes)); // very inefficient, to big
828: PetscCall(ISGetIndices(isrow, &rowindices));
829: PetscCall(ISGetIndices(isicol, &icolindices));
830: nDMs = 0;
831: bsrt = 0;
832: bend = 1;
833: for (PetscInt row_B = 0; row_B < nloc; row_B++) { // for all rows in block diagonal space
834: PetscInt rowA = icolindices[row_B], minj = PETSC_INT_MAX, maxj = 0;
835: //PetscCall(PetscPrintf(PETSC_COMM_SELF, "\t[%d] rowA = %d\n",rank,rowA));
836: PetscCall(MatGetRow(Aseq, rowA, &ncols, &colsA, NULL)); // not sorted in permutation
837: PetscCheck(ncols, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Empty row not supported: %" PetscInt_FMT, row_B);
838: for (PetscInt colj = 0; colj < ncols; colj++) {
839: PetscInt colB = rowindices[colsA[colj]]; // use local idx
840: //PetscCall(PetscPrintf(PETSC_COMM_SELF, "\t\t[%d] colB = %d\n",rank,colB));
841: PetscCheck(colB >= 0 && colB < nloc, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "colB < 0: %" PetscInt_FMT, colB);
842: if (colB > maxj) maxj = colB;
843: if (colB < minj) minj = colB;
844: }
845: PetscCall(MatRestoreRow(Aseq, rowA, &ncols, &colsA, NULL));
846: if (minj >= bend) { // first column is > max of last block -- new block or last block
847: //PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\t\t finish block %d, N loc = %d (%d,%d)\n", nDMs+1, bend - bsrt,bsrt,bend));
848: block_sizes[nDMs] = bend - bsrt;
849: ntot += block_sizes[nDMs];
850: PetscCheck(minj == bend, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "minj != bend: %" PetscInt_FMT " != %" PetscInt_FMT, minj, bend);
851: bsrt = bend;
852: bend++; // start with size 1 in new block
853: nDMs++;
854: }
855: if (maxj + 1 > bend) bend = maxj + 1;
856: PetscCheck(minj >= bsrt || row_B == Iend - 1, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "%" PetscInt_FMT ") minj < bsrt: %" PetscInt_FMT " != %" PetscInt_FMT, rowA, minj, bsrt);
857: //PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] %d) row %d.%d) cols %d : %d ; bsrt = %d, bend = %d\n",rank,row_B,nDMs,rowA,minj,maxj,bsrt,bend));
858: }
859: // do last block
860: //PetscCall(PetscPrintf(PETSC_COMM_SELF, "\t\t\t [%d] finish block %d, N loc = %d (%d,%d)\n", rank, nDMs+1, bend - bsrt,bsrt,bend));
861: block_sizes[nDMs] = bend - bsrt;
862: ntot += block_sizes[nDMs];
863: nDMs++;
864: // cleanup
865: PetscCheck(ntot == nloc, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "n total != n local: %" PetscInt_FMT " != %" PetscInt_FMT, ntot, nloc);
866: PetscCall(ISRestoreIndices(isrow, &rowindices));
867: PetscCall(ISRestoreIndices(isicol, &icolindices));
868: PetscCall(PetscRealloc(sizeof(PetscInt) * nDMs, &block_sizes));
869: PetscCall(MatCreateVecs(A, &jac->vec_diag, NULL));
870: PetscCall(PetscInfo(pc, "Setup Matrix based meta data (not DMComposite not attached to PC) %" PetscInt_FMT " sub domains\n", nDMs));
871: }
872: PetscCall(ISDestroy(&isrow));
873: PetscCall(ISDestroy(&isicol));
874: jac->num_dms = nDMs;
875: PetscCall(VecGetLocalSize(jac->vec_diag, &n));
876: jac->n = n;
877: jac->d_idiag_k = new Kokkos::View<PetscScalar *, Kokkos::LayoutRight>("idiag", n);
878: // options
879: PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc));
880: PetscCall(KSPSetFromOptions(jac->ksp));
881: PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPBICG, ""));
882: if (flg) {
883: jac->ksp_type_idx = BATCH_KSP_BICG_IDX;
884: jac->nwork = 7;
885: } else {
886: PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPTFQMR, ""));
887: if (flg) {
888: jac->ksp_type_idx = BATCH_KSP_TFQMR_IDX;
889: jac->nwork = 10;
890: } else {
891: #if PetscDefined(HAVE_KOKKOS_KERNELS_BATCH)
892: PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPGMRES, ""));
893: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Unsupported batch ksp type");
894: jac->ksp_type_idx = BATCH_KSP_GMRESKK_IDX;
895: jac->nwork = 0;
896: #else
897: KSPType ksptype;
898: PetscCall(KSPGetType(jac->ksp, &ksptype));
899: PetscCheck(flg, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Type: %s not supported in complex", ksptype);
900: #endif
901: }
902: }
903: PetscOptionsBegin(PetscObjectComm((PetscObject)jac->ksp), ((PetscObject)jac->ksp)->prefix, "Options for Kokkos batch solver", "none");
904: PetscCall(PetscOptionsBool("-ksp_converged_reason", "", "bjkokkos.kokkos.cxx.c", jac->reason, &jac->reason, NULL));
905: PetscCall(PetscOptionsBool("-ksp_monitor", "", "bjkokkos.kokkos.cxx.c", jac->monitor, &jac->monitor, NULL));
906: PetscCall(PetscOptionsInt("-ksp_batch_target", "", "bjkokkos.kokkos.cxx.c", jac->batch_target, &jac->batch_target, NULL));
907: PetscCall(PetscOptionsInt("-ksp_rank_target", "", "bjkokkos.kokkos.cxx.c", jac->rank_target, &jac->rank_target, NULL));
908: PetscCall(PetscOptionsInt("-ksp_batch_nsolves_team", "", "bjkokkos.kokkos.cxx.c", jac->nsolves_team, &jac->nsolves_team, NULL));
909: PetscCheck(jac->batch_target < jac->num_dms, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "-ksp_batch_target (%" PetscInt_FMT ") >= number of DMs (%" PetscInt_FMT ")", jac->batch_target, jac->num_dms);
910: PetscOptionsEnd();
911: // get blocks - jac->d_bid_eqOffset_k
912: if (pack) {
913: PetscCall(PetscMalloc(sizeof(*subX) * nDMs, &subX));
914: PetscCall(PetscMalloc(sizeof(*subDM) * nDMs, &subDM));
915: }
916: PetscCall(PetscMalloc(sizeof(*jac->dm_Nf) * nDMs, &jac->dm_Nf));
917: PetscCall(PetscInfo(pc, "Have %" PetscInt_FMT " blocks, n=%" PetscInt_FMT " rtol=%g type = %s\n", nDMs, n, (double)jac->ksp->rtol, ((PetscObject)jac->ksp)->type_name));
918: if (pack) PetscCall(DMCompositeGetEntriesArray(pack, subDM));
919: jac->nBlocks = 0;
920: for (PetscInt ii = 0; ii < nDMs; ii++) {
921: PetscInt Nf;
922: if (subDM) {
923: DM dm = subDM[ii];
924: PetscSection section;
925: PetscCall(DMGetLocalSection(dm, §ion));
926: PetscCall(PetscSectionGetNumFields(section, &Nf));
927: } else Nf = 1;
928: jac->nBlocks += Nf;
929: #if PCBJKOKKOS_VERBOSE_LEVEL <= 2
930: if (ii == 0) PetscCall(PetscInfo(pc, "%" PetscInt_FMT ") %" PetscInt_FMT " blocks (%" PetscInt_FMT " total)\n", ii, Nf, jac->nBlocks));
931: #else
932: PetscCall(PetscInfo(pc, "%" PetscInt_FMT ") %" PetscInt_FMT " blocks (%" PetscInt_FMT " total)\n", ii, Nf, jac->nBlocks));
933: #endif
934: jac->dm_Nf[ii] = Nf;
935: }
936: { // d_bid_eqOffset_k
937: Kokkos::View<PetscInt *, Kokkos::LayoutRight, Kokkos::HostSpace> h_block_offsets("block_offsets", jac->nBlocks + 1);
938: if (pack) PetscCall(DMCompositeGetAccessArray(pack, jac->vec_diag, nDMs, NULL, subX));
939: h_block_offsets[0] = 0;
940: jac->const_block_size = -1;
941: for (PetscInt ii = 0, idx = 0; ii < nDMs; ii++) {
942: PetscInt nloc, nblk;
943: if (pack) PetscCall(VecGetSize(subX[ii], &nloc));
944: else nloc = block_sizes[ii];
945: nblk = nloc / jac->dm_Nf[ii];
946: PetscCheck(nloc % jac->dm_Nf[ii] == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_USER, "nloc%%jac->dm_Nf[ii] (%" PetscInt_FMT ") != 0 DMs", nloc % jac->dm_Nf[ii]);
947: for (PetscInt jj = 0; jj < jac->dm_Nf[ii]; jj++, idx++) {
948: h_block_offsets[idx + 1] = h_block_offsets[idx] + nblk;
949: #if PCBJKOKKOS_VERBOSE_LEVEL <= 2
950: if (idx == 0) PetscCall(PetscInfo(pc, "Add first of %" PetscInt_FMT " blocks with %" PetscInt_FMT " equations\n", jac->nBlocks, nblk));
951: #else
952: PetscCall(PetscInfo(pc, "\t%" PetscInt_FMT ") Add block with %" PetscInt_FMT " equations of %" PetscInt_FMT "\n", idx + 1, nblk, jac->nBlocks));
953: #endif
954: if (jac->const_block_size == -1) jac->const_block_size = nblk;
955: else if (jac->const_block_size > 0 && jac->const_block_size != nblk) jac->const_block_size = 0;
956: }
957: }
958: if (pack) {
959: PetscCall(DMCompositeRestoreAccessArray(pack, jac->vec_diag, jac->nBlocks, NULL, subX));
960: PetscCall(PetscFree(subX));
961: PetscCall(PetscFree(subDM));
962: }
963: jac->d_bid_eqOffset_k = new Kokkos::View<PetscInt *, Kokkos::LayoutRight>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(), h_block_offsets));
964: Kokkos::deep_copy(*jac->d_bid_eqOffset_k, h_block_offsets);
965: }
966: if (!pack) PetscCall(PetscFree(block_sizes));
967: }
968: { // get jac->d_idiag_k (PC setup),
969: const PetscInt *d_ai, *d_aj;
970: const PetscScalar *d_aa;
971: const PetscInt conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0 && PCBJKOKKOS_VEC_SIZE != 1) ? PCBJKOKKOS_TEAM_SIZE : 1;
972: const PetscInt *d_bid_eqOffset = jac->d_bid_eqOffset_k->data(), *r = jac->d_isrow_k->data(), *ic = jac->d_isicol_k->data();
973: PetscScalar *d_idiag = jac->d_idiag_k->data(), *dummy;
974: PetscMemType mtype;
975: PetscCall(MatSeqAIJGetCSRAndMemType(Aseq, &d_ai, &d_aj, &dummy, &mtype));
976: d_aa = dummy;
977: Kokkos::parallel_for(
978: "Diag", Kokkos::TeamPolicy<>(jac->nBlocks, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) {
979: const PetscInt blkID = team.league_rank();
980: Kokkos::parallel_for(Kokkos::TeamThreadRange(team, d_bid_eqOffset[blkID], d_bid_eqOffset[blkID + 1]), [=](const int rowb) {
981: const PetscInt rowa = ic[rowb], ai = d_ai[rowa], *aj = d_aj + ai; // grab original data
982: const PetscScalar *aa = d_aa + ai;
983: const PetscInt nrow = d_ai[rowa + 1] - ai;
984: int found;
985: Kokkos::parallel_reduce(
986: Kokkos::ThreadVectorRange(team, nrow),
987: [=](const int &j, int &count) {
988: const PetscInt colb = r[aj[j]];
989: if (colb == rowb) {
990: d_idiag[rowb] = 1. / aa[j];
991: count++;
992: }
993: },
994: found);
995: #if PetscDefined(USE_DEBUG) && !PetscDefined(HAVE_SYCL)
996: if (found != 1) Kokkos::single(Kokkos::PerThread(team), [=]() { printf("ERRORrow %d) found = %d\n", rowb, found); });
997: #endif
998: });
999: });
1000: }
1001: }
1002: PetscFunctionReturn(PETSC_SUCCESS);
1003: }
1005: /* Default destroy, if it has never been setup */
1006: static PetscErrorCode PCReset_BJKOKKOS(PC pc)
1007: {
1008: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
1010: PetscFunctionBegin;
1011: PetscCall(KSPDestroy(&jac->ksp));
1012: PetscCall(VecDestroy(&jac->vec_diag));
1013: if (jac->d_bid_eqOffset_k) delete jac->d_bid_eqOffset_k;
1014: if (jac->d_idiag_k) delete jac->d_idiag_k;
1015: if (jac->d_isrow_k) delete jac->d_isrow_k;
1016: if (jac->d_isicol_k) delete jac->d_isicol_k;
1017: jac->d_bid_eqOffset_k = NULL;
1018: jac->d_idiag_k = NULL;
1019: jac->d_isrow_k = NULL;
1020: jac->d_isicol_k = NULL;
1021: PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSGetKSP_C", NULL)); // not published now (causes configure errors)
1022: PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSSetKSP_C", NULL));
1023: PetscCall(PetscFree(jac->dm_Nf));
1024: jac->dm_Nf = NULL;
1025: if (jac->rowOffsets) delete jac->rowOffsets;
1026: if (jac->colIndices) delete jac->colIndices;
1027: if (jac->batch_b) delete jac->batch_b;
1028: if (jac->batch_x) delete jac->batch_x;
1029: if (jac->batch_values) delete jac->batch_values;
1030: jac->rowOffsets = NULL;
1031: jac->colIndices = NULL;
1032: jac->batch_b = NULL;
1033: jac->batch_x = NULL;
1034: jac->batch_values = NULL;
1035: PetscFunctionReturn(PETSC_SUCCESS);
1036: }
1038: static PetscErrorCode PCDestroy_BJKOKKOS(PC pc)
1039: {
1040: PetscFunctionBegin;
1041: PetscCall(PCReset_BJKOKKOS(pc));
1042: PetscCall(PetscFree(pc->data));
1043: PetscFunctionReturn(PETSC_SUCCESS);
1044: }
1046: static PetscErrorCode PCView_BJKOKKOS(PC pc, PetscViewer viewer)
1047: {
1048: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
1049: PetscBool isascii;
1051: PetscFunctionBegin;
1052: if (!jac->ksp) PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc));
1053: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
1054: if (isascii) {
1055: PetscCall(PetscViewerASCIIPrintf(viewer, " Batched device linear solver: Krylov (KSP) method with Jacobi preconditioning\n"));
1056: PetscCall(PetscViewerASCIIPrintf(viewer, "\t\tnwork = %" PetscInt_FMT ", rel tol = %e, abs tol = %e, div tol = %e, max it =%" PetscInt_FMT ", type = %s\n", jac->nwork, jac->ksp->rtol, jac->ksp->abstol, jac->ksp->divtol, jac->ksp->max_it,
1057: ((PetscObject)jac->ksp)->type_name));
1058: }
1059: PetscFunctionReturn(PETSC_SUCCESS);
1060: }
1062: static PetscErrorCode PCSetFromOptions_BJKOKKOS(PC pc, PetscOptionItems PetscOptionsObject)
1063: {
1064: PetscFunctionBegin;
1065: PetscOptionsHeadBegin(PetscOptionsObject, "PC BJKOKKOS options");
1066: PetscOptionsHeadEnd();
1067: PetscFunctionReturn(PETSC_SUCCESS);
1068: }
1070: static PetscErrorCode PCBJKOKKOSSetKSP_BJKOKKOS(PC pc, KSP ksp)
1071: {
1072: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
1074: PetscFunctionBegin;
1075: PetscCall(PetscObjectReference((PetscObject)ksp));
1076: PetscCall(KSPDestroy(&jac->ksp));
1077: jac->ksp = ksp;
1078: PetscFunctionReturn(PETSC_SUCCESS);
1079: }
1081: /*@
1082: PCBJKOKKOSSetKSP - Sets the `KSP` context for `PCBJKOKKOS`
1084: Collective
1086: Input Parameters:
1087: + pc - the `PCBJKOKKOS` preconditioner context
1088: - ksp - the `KSP` solver
1090: Level: advanced
1092: Notes:
1093: The `PC` and the `KSP` must have the same communicator
1095: If the `PC` is not `PCBJKOKKOS` this function returns without doing anything
1097: .seealso: [](ch_ksp), `PCBJKOKKOSGetKSP()`, `PCBJKOKKOS`
1098: @*/
1099: PetscErrorCode PCBJKOKKOSSetKSP(PC pc, KSP ksp)
1100: {
1101: PetscFunctionBegin;
1104: PetscCheckSameComm(pc, 1, ksp, 2);
1105: PetscTryMethod(pc, "PCBJKOKKOSSetKSP_C", (PC, KSP), (pc, ksp));
1106: PetscFunctionReturn(PETSC_SUCCESS);
1107: }
1109: static PetscErrorCode PCBJKOKKOSGetKSP_BJKOKKOS(PC pc, KSP *ksp)
1110: {
1111: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
1113: PetscFunctionBegin;
1114: if (!jac->ksp) PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc));
1115: *ksp = jac->ksp;
1116: PetscFunctionReturn(PETSC_SUCCESS);
1117: }
1119: /*@
1120: PCBJKOKKOSGetKSP - Gets the `KSP` context for the `PCBJKOKKOS` preconditioner
1122: Not Collective but `KSP` returned is parallel if `PC` was parallel
1124: Input Parameter:
1125: . pc - the preconditioner context
1127: Output Parameter:
1128: . ksp - the `KSP` solver
1130: Level: advanced
1132: Notes:
1133: You must call `KSPSetUp()` before calling `PCBJKOKKOSGetKSP()`.
1135: If the `PC` is not a `PCBJKOKKOS` object it raises an error
1137: .seealso: [](ch_ksp), `PCBJKOKKOS`, `PCBJKOKKOSSetKSP()`
1138: @*/
1139: PetscErrorCode PCBJKOKKOSGetKSP(PC pc, KSP *ksp)
1140: {
1141: PetscFunctionBegin;
1143: PetscAssertPointer(ksp, 2);
1144: PetscUseMethod(pc, "PCBJKOKKOSGetKSP_C", (PC, KSP *), (pc, ksp));
1145: PetscFunctionReturn(PETSC_SUCCESS);
1146: }
1148: static PetscErrorCode PCPostSolve_BJKOKKOS(PC pc, KSP ksp, Vec b, Vec x)
1149: {
1150: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
1152: PetscFunctionBegin;
1154: ksp->its = jac->max_nits;
1155: PetscFunctionReturn(PETSC_SUCCESS);
1156: }
1158: static PetscErrorCode PCPreSolve_BJKOKKOS(PC pc, KSP ksp, Vec b, Vec x)
1159: {
1160: PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
1162: PetscFunctionBegin;
1164: jac->ksp->errorifnotconverged = ksp->errorifnotconverged;
1165: PetscFunctionReturn(PETSC_SUCCESS);
1166: }
1168: /*MC
1169: PCBJKOKKOS - A batched Krylov/block Jacobi solver that runs a solve of each diagaonl block of a block diagonal `MATSEQAIJ` in a Kokkos thread group
1171: Options Database Key:
1172: . -pc_bjkokkos_ - options prefix for its `KSP` options
1174: Level: intermediate
1176: Note:
1177: For use with `-ksp_type preonly` to bypass any computation on the CPU
1179: Developer Notes:
1180: The entire Krylov (TFQMR or BICG) with diagonal preconditioning for each block of a block diagnaol matrix runs in a Kokkos thread group (eg, one block per SM on NVIDIA). It supports taking a non-block diagonal matrix but this is not tested. One should create an explicit block diagonal matrix and use that as the matrix for constructing the preconditioner in the outer `KSP` solver. Variable block size are supported and tested in src/ts/utils/dmplexlandau/tutorials/ex[1|2].c
1182: .seealso: [](ch_ksp), `PCCreate()`, `PCSetType()`, `PCType`, `PC`, `PCBJACOBI`,
1183: `PCSHELL`, `PCCOMPOSITE`, `PCSetUseAmat()`, `PCBJKOKKOSGetKSP()`
1184: M*/
1186: PETSC_EXTERN PetscErrorCode PCCreate_BJKOKKOS(PC pc)
1187: {
1188: PC_PCBJKOKKOS *jac;
1190: PetscFunctionBegin;
1191: PetscCall(PetscNew(&jac));
1192: pc->data = (void *)jac;
1194: jac->ksp = NULL;
1195: jac->vec_diag = NULL;
1196: jac->d_bid_eqOffset_k = NULL;
1197: jac->d_idiag_k = NULL;
1198: jac->d_isrow_k = NULL;
1199: jac->d_isicol_k = NULL;
1200: jac->nBlocks = 1;
1201: jac->max_nits = 0;
1203: PetscCall(PetscMemzero(pc->ops, sizeof(struct _PCOps)));
1204: pc->ops->apply = PCApply_BJKOKKOS;
1205: pc->ops->applytranspose = NULL;
1206: pc->ops->setup = PCSetUp_BJKOKKOS;
1207: pc->ops->reset = PCReset_BJKOKKOS;
1208: pc->ops->destroy = PCDestroy_BJKOKKOS;
1209: pc->ops->setfromoptions = PCSetFromOptions_BJKOKKOS;
1210: pc->ops->view = PCView_BJKOKKOS;
1211: pc->ops->postsolve = PCPostSolve_BJKOKKOS;
1212: pc->ops->presolve = PCPreSolve_BJKOKKOS;
1214: jac->rowOffsets = NULL;
1215: jac->colIndices = NULL;
1216: jac->batch_b = NULL;
1217: jac->batch_x = NULL;
1218: jac->batch_values = NULL;
1220: PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSGetKSP_C", PCBJKOKKOSGetKSP_BJKOKKOS));
1221: PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSSetKSP_C", PCBJKOKKOSSetKSP_BJKOKKOS));
1222: PetscFunctionReturn(PETSC_SUCCESS);
1223: }