Date: 2025-03-28T21:11:02+00:00

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: #if defined(PETSC_USE_COMPLEX)
     75:   const int aconj = A->hermitian == PETSC_BOOL3_TRUE;
     76: #else
     77:   const int aconj = 0;
     78: #endif

     80:   PetscFunctionBegin;
     81:   PetscCall(VecSet(zz, 0.0));
     82:   PetscCall(VecGetArrayRead(xx, &x));
     83:   PetscCall(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:   PetscCall(VecRestoreArrayRead(xx, &x));
    120:   PetscCall(VecRestoreArray(zz, &z));
    121:   PetscCall(PetscLogFlops(2.0 * (2.0 * a->nz - nonzerorow) - nonzerorow));
    122:   PetscFunctionReturn(PETSC_SUCCESS);
    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:   PetscFunctionBegin;
    146:   if (fshift == -1.0) fshift = 0.0; /* negative fshift indicates do not error on zero diagonal; this code never errors on zero diagonal */
    147:   PetscCheck(!(flag & SOR_EISENSTAT), PETSC_COMM_SELF, PETSC_ERR_SUP, "No support yet for Eisenstat");

    149:   its = its * lits;
    150:   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);

    152:   PetscCheck(bs <= 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "SSOR for block size > 1 is not yet implemented");

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

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

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

    166:   aidiag = a->idiag;

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

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

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

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

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

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

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

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

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

    311:   PetscCall(VecRestoreArray(xx, &x));
    312:   PetscCall(VecRestoreArrayRead(bb, &b));
    313:   PetscFunctionReturn(PETSC_SUCCESS);
    314: }