Actual source code: relax.h


  2: /*
  3:     This is included by sbaij.c to generate unsigned short and regular versions of these two functions
  4: */

  6: /* We cut-and-past below from aij.h to make a "no_function" version of PetscSparseDensePlusDot().
  7:  * This is necessary because the USESHORT case cannot use the inlined functions that may be employed. */

  9: #if defined(PETSC_KERNEL_USE_UNROLL_4)
 10:   #define PetscSparseDensePlusDot_no_function(sum, r, xv, xi, nnz) \
 11:     { \
 12:       if (nnz > 0) { \
 13:         PetscInt nnz2 = nnz, rem = nnz & 0x3; \
 14:         switch (rem) { \
 15:         case 3: \
 16:           sum += *xv++ * r[*xi++]; \
 17:         case 2: \
 18:           sum += *xv++ * r[*xi++]; \
 19:         case 1: \
 20:           sum += *xv++ * r[*xi++]; \
 21:           nnz2 -= rem; \
 22:         } \
 23:         while (nnz2 > 0) { \
 24:           sum += xv[0] * r[xi[0]] + xv[1] * r[xi[1]] + xv[2] * r[xi[2]] + xv[3] * r[xi[3]]; \
 25:           xv += 4; \
 26:           xi += 4; \
 27:           nnz2 -= 4; \
 28:         } \
 29:         xv -= nnz; \
 30:         xi -= nnz; \
 31:       } \
 32:     }

 34: #elif defined(PETSC_KERNEL_USE_UNROLL_2)
 35:   #define PetscSparseDensePlusDot_no_function(sum, r, xv, xi, nnz) \
 36:     { \
 37:       PetscInt __i, __i1, __i2; \
 38:       for (__i = 0; __i < nnz - 1; __i += 2) { \
 39:         __i1 = xi[__i]; \
 40:         __i2 = xi[__i + 1]; \
 41:         sum += (xv[__i] * r[__i1] + xv[__i + 1] * r[__i2]); \
 42:       } \
 43:       if (nnz & 0x1) sum += xv[__i] * r[xi[__i]]; \
 44:     }

 46: #else
 47:   #define PetscSparseDensePlusDot_no_function(sum, r, xv, xi, nnz) \
 48:     { \
 49:       PetscInt __i; \
 50:       for (__i = 0; __i < nnz; __i++) sum += xv[__i] * r[xi[__i]]; \
 51:     }
 52: #endif

 54: #if defined(USESHORT)
 55: PetscErrorCode MatMult_SeqSBAIJ_1_ushort(Mat A, Vec xx, Vec zz)
 56: #else
 57: PetscErrorCode MatMult_SeqSBAIJ_1(Mat A, Vec xx, Vec zz)
 58: #endif
 59: {
 60:   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ *)A->data;
 61:   const PetscScalar *x;
 62:   PetscScalar       *z, x1, sum;
 63:   const MatScalar   *v;
 64:   MatScalar          vj;
 65:   PetscInt           mbs = a->mbs, i, j, nz;
 66:   const PetscInt    *ai  = a->i;
 67: #if defined(USESHORT)
 68:   const unsigned short *ib = a->jshort;
 69:   unsigned short        ibt;
 70: #else
 71:   const PetscInt *ib = a->j;
 72:   PetscInt        ibt;
 73: #endif
 74:   PetscInt nonzerorow = 0, jmin;
 75: #if defined(PETSC_USE_COMPLEX)
 76:   const int aconj = A->hermitian == PETSC_BOOL3_TRUE;
 77: #else
 78:   const int       aconj = 0;
 79: #endif

 81:   VecSet(zz, 0.0);
 82:   VecGetArrayRead(xx, &x);
 83:   VecGetArray(zz, &z);

 85:   v = a->a;
 86:   for (i = 0; i < mbs; i++) {
 87:     nz = ai[i + 1] - ai[i]; /* length of i_th row of A */
 88:     if (!nz) continue;      /* Move to the next row if the current row is empty */
 89:     nonzerorow++;
 90:     sum  = 0.0;
 91:     jmin = 0;
 92:     x1   = x[i];
 93:     if (ib[0] == i) {
 94:       sum = v[0] * x1; /* diagonal term */
 95:       jmin++;
 96:     }
 97:     PetscPrefetchBlock(ib + nz, nz, 0, PETSC_PREFETCH_HINT_NTA); /* Indices for the next row (assumes same size as this one) */
 98:     PetscPrefetchBlock(v + nz, nz, 0, PETSC_PREFETCH_HINT_NTA);  /* Entries for the next row */
 99:     if (aconj) {
100:       for (j = jmin; j < nz; j++) {
101:         ibt = ib[j];
102:         vj  = v[j];
103:         z[ibt] += PetscConj(vj) * x1; /* (strict lower triangular part of A)*x  */
104:         sum += vj * x[ibt];           /* (strict upper triangular part of A)*x  */
105:       }
106:     } else {
107:       for (j = jmin; j < nz; j++) {
108:         ibt = ib[j];
109:         vj  = v[j];
110:         z[ibt] += vj * x1;  /* (strict lower triangular part of A)*x  */
111:         sum += vj * x[ibt]; /* (strict upper triangular part of A)*x  */
112:       }
113:     }
114:     z[i] += sum;
115:     v += nz;
116:     ib += nz;
117:   }

119:   VecRestoreArrayRead(xx, &x);
120:   VecRestoreArray(zz, &z);
121:   PetscLogFlops(2.0 * (2.0 * a->nz - nonzerorow) - nonzerorow);
122:   return 0;
123: }

125: #if defined(USESHORT)
126: PetscErrorCode MatSOR_SeqSBAIJ_ushort(Mat A, Vec bb, PetscReal omega, MatSORType flag, PetscReal fshift, PetscInt its, PetscInt lits, Vec xx)
127: #else
128: PetscErrorCode MatSOR_SeqSBAIJ(Mat A, Vec bb, PetscReal omega, MatSORType flag, PetscReal fshift, PetscInt its, PetscInt lits, Vec xx)
129: #endif
130: {
131:   Mat_SeqSBAIJ      *a  = (Mat_SeqSBAIJ *)A->data;
132:   const MatScalar   *aa = a->a, *v, *v1, *aidiag;
133:   PetscScalar       *x, *t, sum;
134:   const PetscScalar *b;
135:   MatScalar          tmp;
136:   PetscInt           m = a->mbs, bs = A->rmap->bs, j;
137:   const PetscInt    *ai = a->i;
138: #if defined(USESHORT)
139:   const unsigned short *aj = a->jshort, *vj, *vj1;
140: #else
141:   const PetscInt *aj = a->j, *vj, *vj1;
142: #endif
143:   PetscInt nz, nz1, i;

145:   if (fshift == -1.0) fshift = 0.0; /* negative fshift indicates do not error on zero diagonal; this code never errors on zero diagonal */

148:   its = its * lits;


153:   VecGetArray(xx, &x);
154:   VecGetArrayRead(bb, &b);

156:   if (!a->idiagvalid) {
157:     if (!a->idiag) PetscMalloc1(m, &a->idiag);
158:     for (i = 0; i < a->mbs; i++) a->idiag[i] = 1.0 / a->a[a->i[i]];
159:     a->idiagvalid = PETSC_TRUE;
160:   }

162:   if (!a->sor_work) PetscMalloc1(m, &a->sor_work);
163:   t = a->sor_work;

165:   aidiag = a->idiag;

167:   if (flag == SOR_APPLY_UPPER) {
168:     /* apply (U + D/omega) to the vector */
169:     PetscScalar d;
170:     for (i = 0; i < m; i++) {
171:       d   = fshift + aa[ai[i]];
172:       nz  = ai[i + 1] - ai[i] - 1;
173:       vj  = aj + ai[i] + 1;
174:       v   = aa + ai[i] + 1;
175:       sum = b[i] * d / omega;
176: #ifdef USESHORT
177:       PetscSparseDensePlusDot_no_function(sum, b, v, vj, nz);
178: #else
179:       PetscSparseDensePlusDot(sum, b, v, vj, nz);
180: #endif
181:       x[i] = sum;
182:     }
183:     PetscLogFlops(a->nz);
184:   }

186:   if (flag & SOR_ZERO_INITIAL_GUESS) {
187:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
188:       PetscArraycpy(t, b, m);

190:       v  = aa + 1;
191:       vj = aj + 1;
192:       for (i = 0; i < m; i++) {
193:         nz  = ai[i + 1] - ai[i] - 1;
194:         tmp = -(x[i] = omega * t[i] * aidiag[i]);
195:         for (j = 0; j < nz; j++) t[vj[j]] += tmp * v[j];
196:         v += nz + 1;
197:         vj += nz + 1;
198:       }
199:       PetscLogFlops(2.0 * a->nz);
200:     }

202:     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
203:       int nz2;
204:       if (!(flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP)) {
205: #if defined(PETSC_USE_BACKWARD_LOOP)
206:         v  = aa + ai[m] - 1;
207:         vj = aj + ai[m] - 1;
208:         for (i = m - 1; i >= 0; i--) {
209:           sum = b[i];
210:           nz  = ai[i + 1] - ai[i] - 1;
211:           {
212:             PetscInt __i;
213:             for (__i = 0; __i < nz; __i++) sum -= v[-__i] * x[vj[-__i]];
214:           }
215: #else
216:         v  = aa + ai[m - 1] + 1;
217:         vj = aj + ai[m - 1] + 1;
218:         nz = 0;
219:         for (i = m - 1; i >= 0; i--) {
220:           sum = b[i];
221:           nz2 = ai[i] - ai[PetscMax(i - 1, 0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
222:           PETSC_Prefetch(v - nz2 - 1, 0, PETSC_PREFETCH_HINT_NTA);
223:           PETSC_Prefetch(vj - nz2 - 1, 0, PETSC_PREFETCH_HINT_NTA);
224:           PetscSparseDenseMinusDot(sum, x, v, vj, nz);
225:           nz = nz2;
226: #endif
227:           x[i] = omega * sum * aidiag[i];
228:           v -= nz + 1;
229:           vj -= nz + 1;
230:         }
231:         PetscLogFlops(2.0 * a->nz);
232:       } else {
233:         v  = aa + ai[m - 1] + 1;
234:         vj = aj + ai[m - 1] + 1;
235:         nz = 0;
236:         for (i = m - 1; i >= 0; i--) {
237:           sum = t[i];
238:           nz2 = ai[i] - ai[PetscMax(i - 1, 0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
239:           PETSC_Prefetch(v - nz2 - 1, 0, PETSC_PREFETCH_HINT_NTA);
240:           PETSC_Prefetch(vj - nz2 - 1, 0, PETSC_PREFETCH_HINT_NTA);
241:           PetscSparseDenseMinusDot(sum, x, v, vj, nz);
242:           x[i] = (1 - omega) * x[i] + omega * sum * aidiag[i];
243:           nz   = nz2;
244:           v -= nz + 1;
245:           vj -= nz + 1;
246:         }
247:         PetscLogFlops(2.0 * a->nz);
248:       }
249:     }
250:     its--;
251:   }

253:   while (its--) {
254:     /*
255:        forward sweep:
256:        for i=0,...,m-1:
257:          sum[i] = (b[i] - U(i,:)x)/d[i];
258:          x[i]   = (1-omega)x[i] + omega*sum[i];
259:          b      = b - x[i]*U^T(i,:);

261:     */
262:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
263:       PetscArraycpy(t, b, m);

265:       for (i = 0; i < m; i++) {
266:         v   = aa + ai[i] + 1;
267:         v1  = v;
268:         vj  = aj + ai[i] + 1;
269:         vj1 = vj;
270:         nz  = ai[i + 1] - ai[i] - 1;
271:         nz1 = nz;
272:         sum = t[i];
273:         while (nz1--) sum -= (*v1++) * x[*vj1++];
274:         x[i] = (1 - omega) * x[i] + omega * sum * aidiag[i];
275:         while (nz--) t[*vj++] -= x[i] * (*v++);
276:       }
277:       PetscLogFlops(4.0 * a->nz);
278:     }

280:     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
281:       /*
282:        backward sweep:
283:        b = b - x[i]*U^T(i,:), i=0,...,n-2
284:        for i=m-1,...,0:
285:          sum[i] = (b[i] - U(i,:)x)/d[i];
286:          x[i]   = (1-omega)x[i] + omega*sum[i];
287:       */
288:       /* if there was a forward sweep done above then I thing the next two for loops are not needed */
289:       PetscArraycpy(t, b, m);

291:       for (i = 0; i < m - 1; i++) { /* update rhs */
292:         v  = aa + ai[i] + 1;
293:         vj = aj + ai[i] + 1;
294:         nz = ai[i + 1] - ai[i] - 1;
295:         while (nz--) t[*vj++] -= x[i] * (*v++);
296:       }
297:       PetscLogFlops(2.0 * (a->nz - m));
298:       for (i = m - 1; i >= 0; i--) {
299:         v   = aa + ai[i] + 1;
300:         vj  = aj + ai[i] + 1;
301:         nz  = ai[i + 1] - ai[i] - 1;
302:         sum = t[i];
303:         while (nz--) sum -= x[*vj++] * (*v++);
304:         x[i] = (1 - omega) * x[i] + omega * sum * aidiag[i];
305:       }
306:       PetscLogFlops(2.0 * (a->nz + m));
307:     }
308:   }

310:   VecRestoreArray(xx, &x);
311:   VecRestoreArrayRead(bb, &b);
312:   return 0;
313: }