xref: /petsc/src/mat/impls/sbaij/seq/relax.h (revision fdf6c4e30aafdbc795e4f76379caa977fd5cdf5a)
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   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 == PETSC_BOOL3_TRUE;
63 #else
64   const int aconj = 0;
65 #endif
66 
67   PetscFunctionBegin;
68   PetscCall(VecSet(zz,0.0));
69   PetscCall(VecGetArrayRead(xx,&x));
70   PetscCall(VecGetArray(zz,&z));
71 
72   v = a->a;
73   for (i=0; i<mbs; i++) {
74     nz = ai[i+1] - ai[i];          /* length of i_th row of A */
75     if (!nz) continue; /* Move to the next row if the current row is empty */
76     nonzerorow++;
77     sum  = 0.0;
78     jmin = 0;
79     x1   = x[i];
80     if (ib[0] == i) {
81       sum = v[0]*x1;                 /* diagonal term */
82       jmin++;
83     }
84     PetscPrefetchBlock(ib+nz,nz,0,PETSC_PREFETCH_HINT_NTA); /* Indices for the next row (assumes same size as this one) */
85     PetscPrefetchBlock(v+nz,nz,0,PETSC_PREFETCH_HINT_NTA);  /* Entries for the next row */
86     if (aconj) {
87       for (j=jmin; j<nz; j++) {
88         ibt     = ib[j];
89         vj      = v[j];
90         z[ibt] += PetscConj(vj) * x1; /* (strict lower triangular part of A)*x  */
91         sum    += vj * x[ibt];        /* (strict upper triangular part of A)*x  */
92       }
93     } else {
94       for (j=jmin; j<nz; j++) {
95         ibt     = ib[j];
96         vj      = v[j];
97         z[ibt] += vj * x1;       /* (strict lower triangular part of A)*x  */
98         sum    += vj * x[ibt];   /* (strict upper triangular part of A)*x  */
99       }
100     }
101     z[i] += sum;
102     v    += nz;
103     ib   += nz;
104   }
105 
106   PetscCall(VecRestoreArrayRead(xx,&x));
107   PetscCall(VecRestoreArray(zz,&z));
108   PetscCall(PetscLogFlops(2.0*(2.0*a->nz - nonzerorow) - nonzerorow));
109   PetscFunctionReturn(0);
110 }
111 
112 #if defined(USESHORT)
113 PetscErrorCode MatSOR_SeqSBAIJ_ushort(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
114 #else
115 PetscErrorCode MatSOR_SeqSBAIJ(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
116 #endif
117 {
118   Mat_SeqSBAIJ      *a = (Mat_SeqSBAIJ*)A->data;
119   const MatScalar   *aa=a->a,*v,*v1,*aidiag;
120   PetscScalar       *x,*t,sum;
121   const PetscScalar *b;
122   MatScalar         tmp;
123   PetscInt          m  =a->mbs,bs=A->rmap->bs,j;
124   const PetscInt    *ai=a->i;
125 #if defined(USESHORT)
126   const unsigned short *aj=a->jshort,*vj,*vj1;
127 #else
128   const PetscInt *aj=a->j,*vj,*vj1;
129 #endif
130   PetscInt nz,nz1,i;
131 
132   PetscFunctionBegin;
133   if (fshift == -1.0) fshift = 0.0; /* negative fshift indicates do not error on zero diagonal; this code never errors on zero diagonal */
134   PetscCheck(!(flag & SOR_EISENSTAT),PETSC_COMM_SELF,PETSC_ERR_SUP,"No support yet for Eisenstat");
135 
136   its = its*lits;
137   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);
138 
139   PetscCheck(bs <= 1,PETSC_COMM_SELF,PETSC_ERR_SUP,"SSOR for block size > 1 is not yet implemented");
140 
141   PetscCall(VecGetArray(xx,&x));
142   PetscCall(VecGetArrayRead(bb,&b));
143 
144   if (!a->idiagvalid) {
145     if (!a->idiag) {
146       PetscCall(PetscMalloc1(m,&a->idiag));
147     }
148     for (i=0; i<a->mbs; i++) a->idiag[i] = 1.0/a->a[a->i[i]];
149     a->idiagvalid = PETSC_TRUE;
150   }
151 
152   if (!a->sor_work) {
153     PetscCall(PetscMalloc1(m,&a->sor_work));
154   }
155   t = a->sor_work;
156 
157   aidiag = a->idiag;
158 
159   if (flag == SOR_APPLY_UPPER) {
160     /* apply (U + D/omega) to the vector */
161     PetscScalar d;
162     for (i=0; i<m; i++) {
163       d   = fshift + aa[ai[i]];
164       nz  = ai[i+1] - ai[i] - 1;
165       vj  = aj + ai[i] + 1;
166       v   = aa + ai[i] + 1;
167       sum = b[i]*d/omega;
168 #ifdef USESHORT
169       PetscSparseDensePlusDot_no_function(sum,b,v,vj,nz);
170 #else
171       PetscSparseDensePlusDot(sum,b,v,vj,nz);
172 #endif
173       x[i] = sum;
174     }
175     PetscCall(PetscLogFlops(a->nz));
176   }
177 
178   if (flag & SOR_ZERO_INITIAL_GUESS) {
179     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
180       PetscCall(PetscArraycpy(t,b,m));
181 
182       v  = aa + 1;
183       vj = aj + 1;
184       for (i=0; i<m; i++) {
185         nz  = ai[i+1] - ai[i] - 1;
186         tmp = -(x[i] = omega*t[i]*aidiag[i]);
187         for (j=0; j<nz; j++) t[vj[j]] += tmp*v[j];
188         v  += nz + 1;
189         vj += nz + 1;
190       }
191       PetscCall(PetscLogFlops(2.0*a->nz));
192     }
193 
194     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
195       int nz2;
196       if (!(flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP)) {
197 #if defined(PETSC_USE_BACKWARD_LOOP)
198         v  = aa + ai[m] - 1;
199         vj = aj + ai[m] - 1;
200         for (i=m-1; i>=0; i--) {
201           sum = b[i];
202           nz  = ai[i+1] - ai[i] - 1;
203           {PetscInt __i;for (__i=0; __i<nz; __i++) sum -= v[-__i] * x[vj[-__i]];}
204 #else
205         v  = aa + ai[m-1] + 1;
206         vj = aj + ai[m-1] + 1;
207         nz = 0;
208         for (i=m-1; i>=0; i--) {
209           sum = b[i];
210           nz2 = ai[i] - ai[PetscMax(i-1,0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
211           PETSC_Prefetch(v-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
212           PETSC_Prefetch(vj-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
213           PetscSparseDenseMinusDot(sum,x,v,vj,nz);
214           nz = nz2;
215 #endif
216           x[i] = omega*sum*aidiag[i];
217           v   -= nz + 1;
218           vj  -= nz + 1;
219         }
220         PetscCall(PetscLogFlops(2.0*a->nz));
221       } else {
222         v  = aa + ai[m-1] + 1;
223         vj = aj + ai[m-1] + 1;
224         nz = 0;
225         for (i=m-1; i>=0; i--) {
226           sum = t[i];
227           nz2 = ai[i] - ai[PetscMax(i-1,0)] - 1; /* avoid referencing ai[-1], nonsense nz2 is okay on last iteration */
228           PETSC_Prefetch(v-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
229           PETSC_Prefetch(vj-nz2-1,0,PETSC_PREFETCH_HINT_NTA);
230           PetscSparseDenseMinusDot(sum,x,v,vj,nz);
231           x[i] = (1-omega)*x[i] + omega*sum*aidiag[i];
232           nz   = nz2;
233           v   -= nz + 1;
234           vj  -= nz + 1;
235         }
236         PetscCall(PetscLogFlops(2.0*a->nz));
237       }
238     }
239     its--;
240   }
241 
242   while (its--) {
243     /*
244        forward sweep:
245        for i=0,...,m-1:
246          sum[i] = (b[i] - U(i,:)x)/d[i];
247          x[i]   = (1-omega)x[i] + omega*sum[i];
248          b      = b - x[i]*U^T(i,:);
249 
250     */
251     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
252       PetscCall(PetscArraycpy(t,b,m));
253 
254       for (i=0; i<m; i++) {
255         v    = aa + ai[i] + 1; v1=v;
256         vj   = aj + ai[i] + 1; vj1=vj;
257         nz   = ai[i+1] - ai[i] - 1; nz1=nz;
258         sum  = t[i];
259         while (nz1--) sum -= (*v1++)*x[*vj1++];
260         x[i] = (1-omega)*x[i] + omega*sum*aidiag[i];
261         while (nz--) t[*vj++] -= x[i]*(*v++);
262       }
263       PetscCall(PetscLogFlops(4.0*a->nz));
264     }
265 
266     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
267       /*
268        backward sweep:
269        b = b - x[i]*U^T(i,:), i=0,...,n-2
270        for i=m-1,...,0:
271          sum[i] = (b[i] - U(i,:)x)/d[i];
272          x[i]   = (1-omega)x[i] + omega*sum[i];
273       */
274       /* if there was a forward sweep done above then I thing the next two for loops are not needed */
275       PetscCall(PetscArraycpy(t,b,m));
276 
277       for (i=0; i<m-1; i++) {  /* update rhs */
278         v    = aa + ai[i] + 1;
279         vj   = aj + ai[i] + 1;
280         nz   = ai[i+1] - ai[i] - 1;
281         while (nz--) t[*vj++] -= x[i]*(*v++);
282       }
283       PetscCall(PetscLogFlops(2.0*(a->nz - m)));
284       for (i=m-1; i>=0; i--) {
285         v    = aa + ai[i] + 1;
286         vj   = aj + ai[i] + 1;
287         nz   = ai[i+1] - ai[i] - 1;
288         sum  = t[i];
289         while (nz--) sum -= x[*vj++]*(*v++);
290         x[i] =   (1-omega)*x[i] + omega*sum*aidiag[i];
291       }
292       PetscCall(PetscLogFlops(2.0*(a->nz + m)));
293     }
294   }
295 
296   PetscCall(VecRestoreArray(xx,&x));
297   PetscCall(VecRestoreArrayRead(bb,&b));
298   PetscFunctionReturn(0);
299 }
300