Actual source code: relax.h
1: /*
2: This is included by sbaij.c to generate unsigned short and regular versions of these two functions
3: */
5: /* We cut-and-past below from aij.h to make a "no_function" version of PetscSparseDensePlusDot().
6: * This is necessary because the USESHORT case cannot use the inlined functions that may be employed. */
8: #if defined(PETSC_KERNEL_USE_UNROLL_4)
9: #define PetscSparseDensePlusDot_no_function(sum, r, xv, xi, nnz) \
10: do { \
11: if (nnz > 0) { \
12: PetscInt nnz2 = nnz, rem = nnz & 0x3; \
13: switch (rem) { \
14: case 3: \
15: sum += *xv++ * r[*xi++]; \
16: case 2: \
17: sum += *xv++ * r[*xi++]; \
18: case 1: \
19: sum += *xv++ * r[*xi++]; \
20: nnz2 -= rem; \
21: } \
22: while (nnz2 > 0) { \
23: sum += xv[0] * r[xi[0]] + xv[1] * r[xi[1]] + xv[2] * r[xi[2]] + xv[3] * r[xi[3]]; \
24: xv += 4; \
25: xi += 4; \
26: nnz2 -= 4; \
27: } \
28: xv -= nnz; \
29: xi -= nnz; \
30: } \
31: } while (0)
33: #elif defined(PETSC_KERNEL_USE_UNROLL_2)
34: #define PetscSparseDensePlusDot_no_function(sum, r, xv, xi, nnz) \
35: do { \
36: PetscInt __i, __i1, __i2; \
37: for (__i = 0; __i < nnz - 1; __i += 2) { \
38: __i1 = xi[__i]; \
39: __i2 = xi[__i + 1]; \
40: sum += (xv[__i] * r[__i1] + xv[__i + 1] * r[__i2]); \
41: } \
42: if (nnz & 0x1) sum += xv[__i] * r[xi[__i]]; \
43: } while (0)
45: #else
46: #define PetscSparseDensePlusDot_no_function(sum, r, xv, xi, nnz) \
47: do { \
48: PetscInt __i; \
49: for (__i = 0; __i < nnz; __i++) sum += xv[__i] * r[xi[__i]]; \
50: } while (0)
51: #endif
53: #if defined(USESHORT)
54: PetscErrorCode MatMult_SeqSBAIJ_1_ushort(Mat A, Vec xx, Vec zz)
55: #else
56: PetscErrorCode MatMult_SeqSBAIJ_1(Mat A, Vec xx, Vec zz)
57: #endif
58: {
59: Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)A->data;
60: const PetscScalar *x;
61: PetscScalar *z, x1, sum;
62: const MatScalar *v;
63: MatScalar vj;
64: PetscInt mbs = a->mbs, i, j, nz;
65: const PetscInt *ai = a->i;
66: #if defined(USESHORT)
67: const unsigned short *ib = a->jshort;
68: unsigned short ibt;
69: #else
70: const PetscInt *ib = a->j;
71: PetscInt ibt;
72: #endif
73: PetscInt nonzerorow = 0, jmin;
74: const int aconj = PetscDefined(USE_COMPLEX) && A->hermitian == PETSC_BOOL3_TRUE ? 1 : 0;
76: PetscFunctionBegin;
77: PetscCall(VecSet(zz, 0.0));
78: PetscCall(VecGetArrayRead(xx, &x));
79: PetscCall(VecGetArray(zz, &z));
81: v = a->a;
82: for (i = 0; i < mbs; i++) {
83: nz = ai[i + 1] - ai[i]; /* length of i_th row of A */
84: if (!nz) continue; /* Move to the next row if the current row is empty */
85: nonzerorow++;
86: sum = 0.0;
87: jmin = 0;
88: x1 = x[i];
89: if (ib[0] == i) {
90: sum = v[0] * x1; /* diagonal term */
91: jmin++;
92: }
93: PetscPrefetchBlock(ib + nz, nz, 0, PETSC_PREFETCH_HINT_NTA); /* Indices for the next row (assumes same size as this one) */
94: PetscPrefetchBlock(v + nz, nz, 0, PETSC_PREFETCH_HINT_NTA); /* Entries for the next row */
95: if (aconj) {
96: for (j = jmin; j < nz; j++) {
97: ibt = ib[j];
98: vj = v[j];
99: z[ibt] += PetscConj(vj) * x1; /* (strict lower triangular part of A)*x */
100: sum += vj * x[ibt]; /* (strict upper triangular part of A)*x */
101: }
102: } else {
103: for (j = jmin; j < nz; j++) {
104: ibt = ib[j];
105: vj = v[j];
106: z[ibt] += vj * x1; /* (strict lower triangular part of A)*x */
107: sum += vj * x[ibt]; /* (strict upper triangular part of A)*x */
108: }
109: }
110: z[i] += sum;
111: v += nz;
112: ib += nz;
113: }
115: PetscCall(VecRestoreArrayRead(xx, &x));
116: PetscCall(VecRestoreArray(zz, &z));
117: PetscCall(PetscLogFlops(2.0 * (2.0 * a->nz - nonzerorow) - nonzerorow));
118: PetscFunctionReturn(PETSC_SUCCESS);
119: }
121: #if defined(USESHORT)
122: PetscErrorCode MatSOR_SeqSBAIJ_ushort(Mat A, Vec bb, PetscReal omega, MatSORType flag, PetscReal fshift, PetscInt its, PetscInt lits, Vec xx)
123: #else
124: PetscErrorCode MatSOR_SeqSBAIJ(Mat A, Vec bb, PetscReal omega, MatSORType flag, PetscReal fshift, PetscInt its, PetscInt lits, Vec xx)
125: #endif
126: {
127: Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)A->data;
128: const MatScalar *aa = a->a, *v, *v1, *aidiag;
129: PetscScalar *x, *t, sum;
130: const PetscScalar *b;
131: MatScalar tmp;
132: PetscInt m = a->mbs, bs = A->rmap->bs, j;
133: const PetscInt *ai = a->i;
134: #if defined(USESHORT)
135: const unsigned short *aj = a->jshort, *vj, *vj1;
136: #else
137: const PetscInt *aj = a->j, *vj, *vj1;
138: #endif
139: PetscInt nz, nz1, i;
141: PetscFunctionBegin;
142: if (fshift == -1.0) fshift = 0.0; /* negative fshift indicates do not error on zero diagonal; this code never errors on zero diagonal */
143: PetscCheck(!(flag & SOR_EISENSTAT), PETSC_COMM_SELF, PETSC_ERR_SUP, "No support yet for Eisenstat");
145: its = its * lits;
146: PetscCheck(its > 0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Relaxation requires global its %" PetscInt_FMT " and local its %" PetscInt_FMT " both positive", its, lits);
148: PetscCheck(bs <= 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "SSOR for block size > 1 is not yet implemented");
150: PetscCall(VecGetArray(xx, &x));
151: PetscCall(VecGetArrayRead(bb, &b));
153: if (!a->idiagvalid) {
154: if (!a->idiag) PetscCall(PetscMalloc1(m, &a->idiag));
155: for (i = 0; i < a->mbs; i++) a->idiag[i] = 1.0 / a->a[a->i[i]];
156: a->idiagvalid = PETSC_TRUE;
157: }
159: if (!a->sor_work) PetscCall(PetscMalloc1(m, &a->sor_work));
160: t = a->sor_work;
162: aidiag = a->idiag;
164: if (flag == SOR_APPLY_UPPER) {
165: /* apply (U + D/omega) to the vector */
166: PetscScalar d;
167: for (i = 0; i < m; i++) {
168: d = fshift + aa[ai[i]];
169: nz = ai[i + 1] - ai[i] - 1;
170: vj = aj + ai[i] + 1;
171: v = aa + ai[i] + 1;
172: sum = b[i] * d / omega;
173: #ifdef USESHORT
174: PetscSparseDensePlusDot_no_function(sum, b, v, vj, nz);
175: #else
176: PetscSparseDensePlusDot(sum, b, v, vj, nz);
177: #endif
178: x[i] = sum;
179: }
180: PetscCall(PetscLogFlops(a->nz));
181: }
183: if (flag & SOR_ZERO_INITIAL_GUESS) {
184: if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
185: PetscCall(PetscArraycpy(t, b, m));
187: v = aa + 1;
188: vj = aj + 1;
189: for (i = 0; i < m; i++) {
190: nz = ai[i + 1] - ai[i] - 1;
191: tmp = -(x[i] = omega * t[i] * aidiag[i]);
192: for (j = 0; j < nz; j++) t[vj[j]] += tmp * v[j];
193: v += nz + 1;
194: vj += nz + 1;
195: }
196: PetscCall(PetscLogFlops(2.0 * a->nz));
197: }
199: if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
200: PetscInt nz2;
201: if (!(flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP)) {
202: v = aa + ai[m] - 1;
203: vj = aj + ai[m] - 1;
204: for (i = m - 1; i >= 0; i--) {
205: sum = b[i];
206: nz = ai[i + 1] - ai[i] - 1;
207: {
208: PetscInt __i;
209: for (__i = 0; __i < nz; __i++) sum -= v[-__i] * x[vj[-__i]];
210: }
211: x[i] = omega * sum * aidiag[i];
212: v -= nz + 1;
213: vj -= nz + 1;
214: }
215: PetscCall(PetscLogFlops(2.0 * a->nz));
216: } else {
217: v = aa + ai[m - 1] + 1;
218: vj = aj + ai[m - 1] + 1;
219: nz = 0;
220: for (i = m - 1; i >= 0; i--) {
221: sum = t[i];
222: nz2 = ai[i] - ai[PetscMax(i - 1, 0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
223: PETSC_Prefetch(v - nz2 - 1, 0, PETSC_PREFETCH_HINT_NTA);
224: PETSC_Prefetch(vj - nz2 - 1, 0, PETSC_PREFETCH_HINT_NTA);
225: PetscSparseDenseMinusDot(sum, x, v, vj, nz);
226: x[i] = (1 - omega) * x[i] + omega * sum * aidiag[i];
227: nz = nz2;
228: v -= nz + 1;
229: vj -= nz + 1;
230: }
231: PetscCall(PetscLogFlops(2.0 * a->nz));
232: }
233: }
234: its--;
235: }
237: while (its--) {
238: /*
239: forward sweep:
240: for i=0,...,m-1:
241: sum[i] = (b[i] - U(i,:)x)/d[i];
242: x[i] = (1-omega)x[i] + omega*sum[i];
243: b = b - x[i]*U^T(i,:);
245: */
246: if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
247: PetscCall(PetscArraycpy(t, b, m));
249: for (i = 0; i < m; i++) {
250: v = aa + ai[i] + 1;
251: v1 = v;
252: vj = aj + ai[i] + 1;
253: vj1 = vj;
254: nz = ai[i + 1] - ai[i] - 1;
255: nz1 = nz;
256: sum = t[i];
257: while (nz1--) sum -= (*v1++) * x[*vj1++];
258: x[i] = (1 - omega) * x[i] + omega * sum * aidiag[i];
259: while (nz--) t[*vj++] -= x[i] * (*v++);
260: }
261: PetscCall(PetscLogFlops(4.0 * a->nz));
262: }
264: if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
265: /*
266: backward sweep:
267: b = b - x[i]*U^T(i,:), i=0,...,n-2
268: for i=m-1,...,0:
269: sum[i] = (b[i] - U(i,:)x)/d[i];
270: x[i] = (1-omega)x[i] + omega*sum[i];
271: */
272: /* if there was a forward sweep done above then I thing the next two for loops are not needed */
273: PetscCall(PetscArraycpy(t, b, m));
275: for (i = 0; i < m - 1; i++) { /* update rhs */
276: v = aa + ai[i] + 1;
277: vj = aj + ai[i] + 1;
278: nz = ai[i + 1] - ai[i] - 1;
279: while (nz--) t[*vj++] -= x[i] * (*v++);
280: }
281: PetscCall(PetscLogFlops(2.0 * (a->nz - m)));
282: for (i = m - 1; i >= 0; i--) {
283: v = aa + ai[i] + 1;
284: vj = aj + ai[i] + 1;
285: nz = ai[i + 1] - ai[i] - 1;
286: sum = t[i];
287: while (nz--) sum -= x[*vj++] * (*v++);
288: x[i] = (1 - omega) * x[i] + omega * sum * aidiag[i];
289: }
290: PetscCall(PetscLogFlops(2.0 * (a->nz + m)));
291: }
292: }
294: PetscCall(VecRestoreArray(xx, &x));
295: PetscCall(VecRestoreArrayRead(bb, &b));
296: PetscFunctionReturn(PETSC_SUCCESS);
297: }