xref: /petsc/src/mat/impls/sbaij/seq/relax.h (revision feff33ee0b5b037fa8f9f294dede656a2f85cc47)
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       switch (nnz & 0x3) { \
13       case 3: sum += *xv++ *r[*xi++]; \
14       case 2: sum += *xv++ *r[*xi++]; \
15       case 1: sum += *xv++ *r[*xi++]; \
16         nnz       -= 4;} \
17       while (nnz > 0) { \
18         sum +=  xv[0] * r[xi[0]] + xv[1] * r[xi[1]] + \
19                xv[2] * r[xi[2]] + xv[3] * r[xi[3]]; \
20         xv += 4; xi += 4; nnz -= 4; }}}
21 
22 #elif defined(PETSC_KERNEL_USE_UNROLL_2)
23 #define PetscSparseDensePlusDot_no_function(sum,r,xv,xi,nnz) { \
24     PetscInt __i,__i1,__i2; \
25     for (__i=0; __i<nnz-1; __i+=2) {__i1 = xi[__i]; __i2=xi[__i+1]; \
26                                     sum += (xv[__i]*r[__i1] + xv[__i+1]*r[__i2]);} \
27     if (nnz & 0x1) sum += xv[__i] * r[xi[__i]];}
28 
29 #else
30 #define PetscSparseDensePlusDot_no_function(sum,r,xv,xi,nnz) { \
31     PetscInt __i; \
32     for (__i=0; __i<nnz; __i++) sum += xv[__i] * r[xi[__i]];}
33 #endif
34 
35 
36 #if defined(USESHORT)
37 PetscErrorCode MatMult_SeqSBAIJ_1_Hermitian_ushort(Mat A,Vec xx,Vec zz)
38 #else
39 PetscErrorCode MatMult_SeqSBAIJ_1_Hermitian(Mat A,Vec xx,Vec zz)
40 #endif
41 {
42   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ*)A->data;
43   const PetscScalar *x;
44   PetscScalar       *z,x1,sum;
45   const MatScalar   *v;
46   MatScalar         vj;
47   PetscErrorCode    ierr;
48   PetscInt          mbs=a->mbs,i,j,nz;
49   const PetscInt    *ai=a->i;
50 #if defined(USESHORT)
51   const unsigned short *ib=a->jshort;
52   unsigned short       ibt;
53 #else
54   const PetscInt *ib=a->j;
55   PetscInt       ibt;
56 #endif
57   PetscInt nonzerorow = 0,jmin;
58 
59   PetscFunctionBegin;
60   ierr = VecSet(zz,0.0);CHKERRQ(ierr);
61   ierr = VecGetArrayRead(xx,&x);CHKERRQ(ierr);
62   ierr = VecGetArray(zz,&z);CHKERRQ(ierr);
63 
64   v = a->a;
65   for (i=0; i<mbs; i++) {
66     nz = ai[i+1] - ai[i];    /* length of i_th row of A */
67     if (!nz) continue; /* Move to the next row if the current row is empty */
68     nonzerorow++;
69     x1   = x[i];
70     sum  = 0.0;
71     jmin = 0;
72     if (ib[0] == i) {
73       sum = v[0]*x1;           /* diagonal term */
74       jmin++;
75     }
76     for (j=jmin; j<nz; j++) {
77       ibt     = ib[j];
78       vj      = v[j];
79       sum    += vj * x[ibt]; /* (strict upper triangular part of A)*x  */
80       z[ibt] += PetscConj(v[j]) * x1;    /* (strict lower triangular part of A)*x  */
81     }
82     z[i] += sum;
83     v    +=    nz;
84     ib   += nz;
85   }
86 
87   ierr = VecRestoreArrayRead(xx,&x);CHKERRQ(ierr);
88   ierr = VecRestoreArray(zz,&z);CHKERRQ(ierr);
89   ierr = PetscLogFlops(2.0*(2.0*a->nz - nonzerorow) - nonzerorow);CHKERRQ(ierr);
90   PetscFunctionReturn(0);
91 }
92 
93 #if defined(USESHORT)
94 PetscErrorCode MatMult_SeqSBAIJ_1_ushort(Mat A,Vec xx,Vec zz)
95 #else
96 PetscErrorCode MatMult_SeqSBAIJ_1(Mat A,Vec xx,Vec zz)
97 #endif
98 {
99   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ*)A->data;
100   const PetscScalar *x;
101   PetscScalar       *z,x1,sum;
102   const MatScalar   *v;
103   MatScalar         vj;
104   PetscErrorCode    ierr;
105   PetscInt          mbs=a->mbs,i,j,nz;
106   const PetscInt    *ai=a->i;
107 #if defined(USESHORT)
108   const unsigned short *ib=a->jshort;
109   unsigned short       ibt;
110 #else
111   const PetscInt *ib=a->j;
112   PetscInt       ibt;
113 #endif
114   PetscInt nonzerorow=0,jmin;
115 
116   PetscFunctionBegin;
117   ierr = VecSet(zz,0.0);CHKERRQ(ierr);
118   ierr = VecGetArrayRead(xx,&x);CHKERRQ(ierr);
119   ierr = VecGetArray(zz,&z);CHKERRQ(ierr);
120 
121   v = a->a;
122   for (i=0; i<mbs; i++) {
123     nz = ai[i+1] - ai[i];          /* length of i_th row of A */
124     if (!nz) continue; /* Move to the next row if the current row is empty */
125     nonzerorow++;
126     sum  = 0.0;
127     jmin = 0;
128     x1   = x[i];
129     if (ib[0] == i) {
130       sum = v[0]*x1;                 /* diagonal term */
131       jmin++;
132     }
133     PetscPrefetchBlock(ib+nz,nz,0,PETSC_PREFETCH_HINT_NTA); /* Indices for the next row (assumes same size as this one) */
134     PetscPrefetchBlock(v+nz,nz,0,PETSC_PREFETCH_HINT_NTA);  /* Entries for the next row */
135     for (j=jmin; j<nz; j++) {
136       ibt     = ib[j];
137       vj      = v[j];
138       z[ibt] += vj * x1;       /* (strict lower triangular part of A)*x  */
139       sum    += vj * x[ibt];   /* (strict upper triangular part of A)*x  */
140     }
141     z[i] += sum;
142     v    += nz;
143     ib   += nz;
144   }
145 
146   ierr = VecRestoreArrayRead(xx,&x);CHKERRQ(ierr);
147   ierr = VecRestoreArray(zz,&z);CHKERRQ(ierr);
148   ierr = PetscLogFlops(2.0*(2.0*a->nz - nonzerorow) - nonzerorow);CHKERRQ(ierr);
149   PetscFunctionReturn(0);
150 }
151 
152 #if defined(USESHORT)
153 PetscErrorCode MatSOR_SeqSBAIJ_ushort(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
154 #else
155 PetscErrorCode MatSOR_SeqSBAIJ(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
156 #endif
157 {
158   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ*)A->data;
159   const MatScalar   *aa=a->a,*v,*v1,*aidiag;
160   PetscScalar       *x,*t,sum;
161   const PetscScalar *b;
162   MatScalar         tmp;
163   PetscErrorCode    ierr;
164   PetscInt          m  =a->mbs,bs=A->rmap->bs,j;
165   const PetscInt    *ai=a->i;
166 #if defined(USESHORT)
167   const unsigned short *aj=a->jshort,*vj,*vj1;
168 #else
169   const PetscInt *aj=a->j,*vj,*vj1;
170 #endif
171   PetscInt nz,nz1,i;
172 
173   PetscFunctionBegin;
174   if (fshift == -1.0) fshift = 0.0; /* negative fshift indicates do not error on zero diagonal; this code never errors on zero diagonal */
175   if (flag & SOR_EISENSTAT) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No support yet for Eisenstat");
176 
177   its = its*lits;
178   if (its <= 0) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Relaxation requires global its %D and local its %D both positive",its,lits);
179 
180   if (bs > 1) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"SSOR for block size > 1 is not yet implemented");
181 
182   ierr = VecGetArray(xx,&x);CHKERRQ(ierr);
183   ierr = VecGetArrayRead(bb,&b);CHKERRQ(ierr);
184 
185   if (!a->idiagvalid) {
186     if (!a->idiag) {
187       ierr = PetscMalloc1(m,&a->idiag);CHKERRQ(ierr);
188     }
189     for (i=0; i<a->mbs; i++) a->idiag[i] = 1.0/a->a[a->i[i]];
190     a->idiagvalid = PETSC_TRUE;
191   }
192 
193   if (!a->sor_work) {
194     ierr = PetscMalloc1(m,&a->sor_work);CHKERRQ(ierr);
195   }
196   t = a->sor_work;
197 
198   aidiag = a->idiag;
199 
200   if (flag == SOR_APPLY_UPPER) {
201     /* apply (U + D/omega) to the vector */
202     PetscScalar d;
203     for (i=0; i<m; i++) {
204       d   = fshift + aa[ai[i]];
205       nz  = ai[i+1] - ai[i] - 1;
206       vj  = aj + ai[i] + 1;
207       v   = aa + ai[i] + 1;
208       sum = b[i]*d/omega;
209 #ifdef USESHORT
210       PetscSparseDensePlusDot_no_function(sum,b,v,vj,nz);
211 #else
212       PetscSparseDensePlusDot(sum,b,v,vj,nz);
213 #endif
214       x[i] = sum;
215     }
216     ierr = PetscLogFlops(a->nz);CHKERRQ(ierr);
217   }
218 
219   if (flag & SOR_ZERO_INITIAL_GUESS) {
220     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
221       ierr = PetscMemcpy(t,b,m*sizeof(PetscScalar));CHKERRQ(ierr);
222 
223       v  = aa + 1;
224       vj = aj + 1;
225       for (i=0; i<m; i++) {
226         nz  = ai[i+1] - ai[i] - 1;
227         tmp = -(x[i] = omega*t[i]*aidiag[i]);
228         for (j=0; j<nz; j++) t[vj[j]] += tmp*v[j];
229         v  += nz + 1;
230         vj += nz + 1;
231       }
232       ierr = PetscLogFlops(2*a->nz);CHKERRQ(ierr);
233     }
234 
235     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
236       int nz2;
237       if (!(flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP)) {
238 #if defined(PETSC_USE_BACKWARD_LOOP)
239         v  = aa + ai[m] - 1;
240         vj = aj + ai[m] - 1;
241         for (i=m-1; i>=0; i--) {
242           sum = b[i];
243           nz  = ai[i+1] - ai[i] - 1;
244           {PetscInt __i;for (__i=0; __i<nz; __i++) sum -= v[-__i] * x[vj[-__i]];}
245 #else
246         v  = aa + ai[m-1] + 1;
247         vj = aj + ai[m-1] + 1;
248         nz = 0;
249         for (i=m-1; i>=0; i--) {
250           sum = b[i];
251           nz2 = ai[i] - ai[PetscMax(i-1,0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
252           PETSC_Prefetch(v-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
253           PETSC_Prefetch(vj-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
254           PetscSparseDenseMinusDot(sum,x,v,vj,nz);
255           nz = nz2;
256 #endif
257           x[i] = omega*sum*aidiag[i];
258           v   -= nz + 1;
259           vj  -= nz + 1;
260         }
261         ierr = PetscLogFlops(2*a->nz);CHKERRQ(ierr);
262       } else {
263         v  = aa + ai[m-1] + 1;
264         vj = aj + ai[m-1] + 1;
265         nz = 0;
266         for (i=m-1; i>=0; i--) {
267           sum = t[i];
268           nz2 = ai[i] - ai[PetscMax(i-1,0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
269           PETSC_Prefetch(v-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
270           PETSC_Prefetch(vj-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
271           PetscSparseDenseMinusDot(sum,x,v,vj,nz);
272           x[i] = (1-omega)*x[i] + omega*sum*aidiag[i];
273           nz   = nz2;
274           v   -= nz + 1;
275           vj  -= nz + 1;
276         }
277         ierr = PetscLogFlops(2*a->nz);CHKERRQ(ierr);
278       }
279     }
280     its--;
281   }
282 
283   while (its--) {
284     /*
285        forward sweep:
286        for i=0,...,m-1:
287          sum[i] = (b[i] - U(i,:)x)/d[i];
288          x[i]   = (1-omega)x[i] + omega*sum[i];
289          b      = b - x[i]*U^T(i,:);
290 
291     */
292     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
293       ierr = PetscMemcpy(t,b,m*sizeof(PetscScalar));CHKERRQ(ierr);
294 
295       for (i=0; i<m; i++) {
296         v    = aa + ai[i] + 1; v1=v;
297         vj   = aj + ai[i] + 1; vj1=vj;
298         nz   = ai[i+1] - ai[i] - 1; nz1=nz;
299         sum  = t[i];
300         while (nz1--) sum -= (*v1++)*x[*vj1++];
301         x[i] = (1-omega)*x[i] + omega*sum*aidiag[i];
302         while (nz--) t[*vj++] -= x[i]*(*v++);
303       }
304       ierr = PetscLogFlops(4.0*a->nz);CHKERRQ(ierr);
305     }
306 
307     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
308       /*
309        backward sweep:
310        b = b - x[i]*U^T(i,:), i=0,...,n-2
311        for i=m-1,...,0:
312          sum[i] = (b[i] - U(i,:)x)/d[i];
313          x[i]   = (1-omega)x[i] + omega*sum[i];
314       */
315       /* if there was a forward sweep done above then I thing the next two for loops are not needed */
316       ierr = PetscMemcpy(t,b,m*sizeof(PetscScalar));CHKERRQ(ierr);
317 
318       for (i=0; i<m-1; i++) {  /* update rhs */
319         v    = aa + ai[i] + 1;
320         vj   = aj + ai[i] + 1;
321         nz   = ai[i+1] - ai[i] - 1;
322         while (nz--) t[*vj++] -= x[i]*(*v++);
323       }
324       ierr = PetscLogFlops(2.0*(a->nz - m));CHKERRQ(ierr);
325       for (i=m-1; i>=0; i--) {
326         v    = aa + ai[i] + 1;
327         vj   = aj + ai[i] + 1;
328         nz   = ai[i+1] - ai[i] - 1;
329         sum  = t[i];
330         while (nz--) sum -= x[*vj++]*(*v++);
331         x[i] =   (1-omega)*x[i] + omega*sum*aidiag[i];
332       }
333       ierr = PetscLogFlops(2.0*(a->nz + m));CHKERRQ(ierr);
334     }
335   }
336 
337   ierr = VecRestoreArray(xx,&x);CHKERRQ(ierr);
338   ierr = VecRestoreArrayRead(bb,&b);CHKERRQ(ierr);
339   PetscFunctionReturn(0);
340 }
341