xref: /petsc/src/mat/impls/sbaij/seq/relax.h (revision 7d5fd1e4d9337468ad3f05b65b7facdcd2dfd2a4)
1 
2 /*
3     This is included by sbaij.c to generate unsigned short and regular versions of these two functions
4 */
5 
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. */
8 
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   }
26 
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]];}
33 
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
39 
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
67 
68   PetscFunctionBegin;
69   ierr = VecSet(zz,0.0);CHKERRQ(ierr);
70   ierr = VecGetArrayRead(xx,&x);CHKERRQ(ierr);
71   ierr = VecGetArray(zz,&z);CHKERRQ(ierr);
72 
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   }
106 
107   ierr = VecRestoreArrayRead(xx,&x);CHKERRQ(ierr);
108   ierr = VecRestoreArray(zz,&z);CHKERRQ(ierr);
109   ierr = PetscLogFlops(2.0*(2.0*a->nz - nonzerorow) - nonzerorow);CHKERRQ(ierr);
110   PetscFunctionReturn(0);
111 }
112 
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;
133 
134   PetscFunctionBegin;
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");
137 
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);
140 
141   if (bs > 1) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"SSOR for block size > 1 is not yet implemented");
142 
143   ierr = VecGetArray(xx,&x);CHKERRQ(ierr);
144   ierr = VecGetArrayRead(bb,&b);CHKERRQ(ierr);
145 
146   if (!a->idiagvalid) {
147     if (!a->idiag) {
148       ierr = PetscMalloc1(m,&a->idiag);CHKERRQ(ierr);
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   }
153 
154   if (!a->sor_work) {
155     ierr = PetscMalloc1(m,&a->sor_work);CHKERRQ(ierr);
156   }
157   t = a->sor_work;
158 
159   aidiag = a->idiag;
160 
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     ierr = PetscLogFlops(a->nz);CHKERRQ(ierr);
178   }
179 
180   if (flag & SOR_ZERO_INITIAL_GUESS) {
181     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
182       ierr = PetscArraycpy(t,b,m);CHKERRQ(ierr);
183 
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       ierr = PetscLogFlops(2.0*a->nz);CHKERRQ(ierr);
194     }
195 
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         ierr = PetscLogFlops(2.0*a->nz);CHKERRQ(ierr);
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         ierr = PetscLogFlops(2.0*a->nz);CHKERRQ(ierr);
239       }
240     }
241     its--;
242   }
243 
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,:);
251 
252     */
253     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
254       ierr = PetscArraycpy(t,b,m);CHKERRQ(ierr);
255 
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       ierr = PetscLogFlops(4.0*a->nz);CHKERRQ(ierr);
266     }
267 
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       ierr = PetscArraycpy(t,b,m);CHKERRQ(ierr);
278 
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       ierr = PetscLogFlops(2.0*(a->nz - m));CHKERRQ(ierr);
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       ierr = PetscLogFlops(2.0*(a->nz + m));CHKERRQ(ierr);
295     }
296   }
297 
298   ierr = VecRestoreArray(xx,&x);CHKERRQ(ierr);
299   ierr = VecRestoreArrayRead(bb,&b);CHKERRQ(ierr);
300   PetscFunctionReturn(0);
301 }
302