Date: 2022-03-04T17:37:53+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:   PetscErrorCode    ierr;
     52:   PetscInt          mbs=a->mbs,i,j,nz;
     53:   const PetscInt    *ai=a->i;
     54: #if defined(USESHORT)
     55:   const unsigned short *ib=a->jshort;
     56:   unsigned short       ibt;
     57: #else
     58:   const PetscInt *ib=a->j;
     59:   PetscInt       ibt;
     60: #endif
     61:   PetscInt nonzerorow=0,jmin;
     62: #if defined(PETSC_USE_COMPLEX)
     63:   const int aconj = A->hermitian;
     64: #else
     65:   const int aconj = 0;
     66: #endif

     69:   VecSet(zz,0.0);
     70:   VecGetArrayRead(xx,&x);
     71:   VecGetArray(zz,&z);

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

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

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

    135:   if (fshift == -1.0) fshift = 0.0; /* negative fshift indicates do not error on zero diagonal; this code never errors on zero diagonal */
    136:   if (flag & SOR_EISENSTAT) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No support yet for Eisenstat");

    138:   its = its*lits;
    139:   if (its <= 0) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Relaxation requires global its %D and local its %D both positive",its,lits);

    141:   if (bs > 1) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"SSOR for block size > 1 is not yet implemented");

    143:   VecGetArray(xx,&x);
    144:   VecGetArrayRead(bb,&b);

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

    154:   if (!a->sor_work) {
    155:     PetscMalloc1(m,&a->sor_work);
    156:   }
    157:   t = a->sor_work;

    159:   aidiag = a->idiag;

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

    180:   if (flag & SOR_ZERO_INITIAL_GUESS) {
    181:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
    182:       PetscArraycpy(t,b,m);

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

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

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

    252:     */
    253:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
    254:       PetscArraycpy(t,b,m);

    256:       for (i=0; i<m; i++) {
    257:         v    = aa + ai[i] + 1; v1=v;
    258:         vj   = aj + ai[i] + 1; vj1=vj;
    259:         nz   = ai[i+1] - ai[i] - 1; nz1=nz;
    260:         sum  = t[i];
    261:         while (nz1--) sum -= (*v1++)*x[*vj1++];
    262:         x[i] = (1-omega)*x[i] + omega*sum*aidiag[i];
    263:         while (nz--) t[*vj++] -= x[i]*(*v++);
    264:       }
    265:       PetscLogFlops(4.0*a->nz);
    266:     }

    268:     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
    269:       /*
    270:        backward sweep:
    271:        b = b - x[i]*U^T(i,:), i=0,...,n-2
    272:        for i=m-1,...,0:
    273:          sum[i] = (b[i] - U(i,:)x)/d[i];
    274:          x[i]   = (1-omega)x[i] + omega*sum[i];
    275:       */
    276:       /* if there was a forward sweep done above then I thing the next two for loops are not needed */
    277:       PetscArraycpy(t,b,m);

    279:       for (i=0; i<m-1; i++) {  /* update rhs */
    280:         v    = aa + ai[i] + 1;
    281:         vj   = aj + ai[i] + 1;
    282:         nz   = ai[i+1] - ai[i] - 1;
    283:         while (nz--) t[*vj++] -= x[i]*(*v++);
    284:       }
    285:       PetscLogFlops(2.0*(a->nz - m));
    286:       for (i=m-1; i>=0; i--) {
    287:         v    = aa + ai[i] + 1;
    288:         vj   = aj + ai[i] + 1;
    289:         nz   = ai[i+1] - ai[i] - 1;
    290:         sum  = t[i];
    291:         while (nz--) sum -= x[*vj++]*(*v++);
    292:         x[i] =   (1-omega)*x[i] + omega*sum*aidiag[i];
    293:       }
    294:       PetscLogFlops(2.0*(a->nz + m));
    295:     }
    296:   }

    298:   VecRestoreArray(xx,&x);
    299:   VecRestoreArrayRead(bb,&b);
    300:   return(0);
    301: }