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