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