Date: 2022-03-31T04:46:55+00:00

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:     if (nnz > 0) { \
     12:       PetscInt nnz2=nnz,rem=nnz&0x3; \
     13:       switch (rem) { \
     14:       case 3: sum += *xv++ *r[*xi++]; \
     15:       case 2: sum += *xv++ *r[*xi++]; \
     16:       case 1: sum += *xv++ *r[*xi++]; \
     17:         nnz2      -= rem;} \
     18:       while (nnz2 > 0) { \
     19:         sum +=  xv[0] * r[xi[0]] + xv[1] * r[xi[1]] + \
     20:                 xv[2] * r[xi[2]] + xv[3] * r[xi[3]]; \
     21:         xv += 4; xi += 4; nnz2 -= 4; \
     22:       } \
     23:       xv -= nnz; xi -= nnz; \
     24:     } \
     25:   }

     27: #elif defined(PETSC_KERNEL_USE_UNROLL_2)
     28: #define PetscSparseDensePlusDot_no_function(sum,r,xv,xi,nnz) { \
     29:     PetscInt __i,__i1,__i2; \
     30:     for (__i=0; __i<nnz-1; __i+=2) {__i1 = xi[__i]; __i2=xi[__i+1]; \
     31:                                     sum += (xv[__i]*r[__i1] + xv[__i+1]*r[__i2]);} \
     32:     if (nnz & 0x1) sum += xv[__i] * r[xi[__i]];}

     34: #else
     35: #define PetscSparseDensePlusDot_no_function(sum,r,xv,xi,nnz) { \
     36:     PetscInt __i; \
     37:     for (__i=0; __i<nnz; __i++) sum += xv[__i] * r[xi[__i]];}
     38: #endif

     40: #if defined(USESHORT)
     41: PetscErrorCode MatMult_SeqSBAIJ_1_ushort(Mat A,Vec xx,Vec zz)
     42: #else
     43: PetscErrorCode MatMult_SeqSBAIJ_1(Mat A,Vec xx,Vec zz)
     44: #endif
     45: {
     46:   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ*)A->data;
     47:   const PetscScalar *x;
     48:   PetscScalar       *z,x1,sum;
     49:   const MatScalar   *v;
     50:   MatScalar         vj;
     51:   PetscInt          mbs=a->mbs,i,j,nz;
     52:   const PetscInt    *ai=a->i;
     53: #if defined(USESHORT)
     54:   const unsigned short *ib=a->jshort;
     55:   unsigned short       ibt;
     56: #else
     57:   const PetscInt *ib=a->j;
     58:   PetscInt       ibt;
     59: #endif
     60:   PetscInt nonzerorow=0,jmin;
     61: #if defined(PETSC_USE_COMPLEX)
     62:   const int aconj = A->hermitian;
     63: #else
     64:   const int aconj = 0;
     65: #endif

     67:   VecSet(zz,0.0);
     68:   VecGetArrayRead(xx,&x);
     69:   VecGetArray(zz,&z);

     71:   v = a->a;
     72:   for (i=0; i<mbs; i++) {
     73:     nz = ai[i+1] - ai[i];          /* length of i_th row of A */
     74:     if (!nz) continue; /* Move to the next row if the current row is empty */
     75:     nonzerorow++;
     76:     sum  = 0.0;
     77:     jmin = 0;
     78:     x1   = x[i];
     79:     if (ib[0] == i) {
     80:       sum = v[0]*x1;                 /* diagonal term */
     81:       jmin++;
     82:     }
     83:     PetscPrefetchBlock(ib+nz,nz,0,PETSC_PREFETCH_HINT_NTA); /* Indices for the next row (assumes same size as this one) */
     84:     PetscPrefetchBlock(v+nz,nz,0,PETSC_PREFETCH_HINT_NTA);  /* Entries for the next row */
     85:     if (aconj) {
     86:       for (j=jmin; j<nz; j++) {
     87:         ibt     = ib[j];
     88:         vj      = v[j];
     89:         z[ibt] += PetscConj(vj) * x1; /* (strict lower triangular part of A)*x  */
     90:         sum    += vj * x[ibt];        /* (strict upper triangular part of A)*x  */
     91:       }
     92:     } else {
     93:       for (j=jmin; j<nz; j++) {
     94:         ibt     = ib[j];
     95:         vj      = v[j];
     96:         z[ibt] += vj * x1;       /* (strict lower triangular part of A)*x  */
     97:         sum    += vj * x[ibt];   /* (strict upper triangular part of A)*x  */
     98:       }
     99:     }
    100:     z[i] += sum;
    101:     v    += nz;
    102:     ib   += nz;
    103:   }

    105:   VecRestoreArrayRead(xx,&x);
    106:   VecRestoreArray(zz,&z);
    107:   PetscLogFlops(2.0*(2.0*a->nz - nonzerorow) - nonzerorow);
    108:   return 0;
    109: }

    111: #if defined(USESHORT)
    112: PetscErrorCode MatSOR_SeqSBAIJ_ushort(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
    113: #else
    114: PetscErrorCode MatSOR_SeqSBAIJ(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
    115: #endif
    116: {
    117:   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ*)A->data;
    118:   const MatScalar   *aa=a->a,*v,*v1,*aidiag;
    119:   PetscScalar       *x,*t,sum;
    120:   const PetscScalar *b;
    121:   MatScalar         tmp;
    122:   PetscInt          m  =a->mbs,bs=A->rmap->bs,j;
    123:   const PetscInt    *ai=a->i;
    124: #if defined(USESHORT)
    125:   const unsigned short *aj=a->jshort,*vj,*vj1;
    126: #else
    127:   const PetscInt *aj=a->j,*vj,*vj1;
    128: #endif
    129:   PetscInt nz,nz1,i;

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

    134:   its = its*lits;

    139:   VecGetArray(xx,&x);
    140:   VecGetArrayRead(bb,&b);

    142:   if (!a->idiagvalid) {
    143:     if (!a->idiag) {
    144:       PetscMalloc1(m,&a->idiag);
    145:     }
    146:     for (i=0; i<a->mbs; i++) a->idiag[i] = 1.0/a->a[a->i[i]];
    147:     a->idiagvalid = PETSC_TRUE;
    148:   }

    150:   if (!a->sor_work) {
    151:     PetscMalloc1(m,&a->sor_work);
    152:   }
    153:   t = a->sor_work;

    155:   aidiag = a->idiag;

    157:   if (flag == SOR_APPLY_UPPER) {
    158:     /* apply (U + D/omega) to the vector */
    159:     PetscScalar d;
    160:     for (i=0; i<m; i++) {
    161:       d   = fshift + aa[ai[i]];
    162:       nz  = ai[i+1] - ai[i] - 1;
    163:       vj  = aj + ai[i] + 1;
    164:       v   = aa + ai[i] + 1;
    165:       sum = b[i]*d/omega;
    166: #ifdef USESHORT
    167:       PetscSparseDensePlusDot_no_function(sum,b,v,vj,nz);
    168: #else
    169:       PetscSparseDensePlusDot(sum,b,v,vj,nz);
    170: #endif
    171:       x[i] = sum;
    172:     }
    173:     PetscLogFlops(a->nz);
    174:   }

    176:   if (flag & SOR_ZERO_INITIAL_GUESS) {
    177:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
    178:       PetscArraycpy(t,b,m);

    180:       v  = aa + 1;
    181:       vj = aj + 1;
    182:       for (i=0; i<m; i++) {
    183:         nz  = ai[i+1] - ai[i] - 1;
    184:         tmp = -(x[i] = omega*t[i]*aidiag[i]);
    185:         for (j=0; j<nz; j++) t[vj[j]] += tmp*v[j];
    186:         v  += nz + 1;
    187:         vj += nz + 1;
    188:       }
    189:       PetscLogFlops(2.0*a->nz);
    190:     }

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

    240:   while (its--) {
    241:     /*
    242:        forward sweep:
    243:        for i=0,...,m-1:
    244:          sum[i] = (b[i] - U(i,:)x)/d[i];
    245:          x[i]   = (1-omega)x[i] + omega*sum[i];
    246:          b      = b - x[i]*U^T(i,:);

    248:     */
    249:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
    250:       PetscArraycpy(t,b,m);

    252:       for (i=0; i<m; i++) {
    253:         v    = aa + ai[i] + 1; v1=v;
    254:         vj   = aj + ai[i] + 1; vj1=vj;
    255:         nz   = ai[i+1] - ai[i] - 1; 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:       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:       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:       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:       PetscLogFlops(2.0*(a->nz + m));
    291:     }
    292:   }

    294:   VecRestoreArray(xx,&x);
    295:   VecRestoreArrayRead(bb,&b);
    296:   return 0;
    297: }