xref: /petsc/src/mat/impls/baij/seq/baijsolvtran6.c (revision ccb4e88a40f0b86eaeca07ff64c64e4de2fae686)
1 #include <../src/mat/impls/baij/seq/baij.h>
2 #include <petsc/private/kernels/blockinvert.h>
3 
4 PetscErrorCode MatSolveTranspose_SeqBAIJ_6_inplace(Mat A,Vec bb,Vec xx)
5 {
6   Mat_SeqBAIJ       *a   =(Mat_SeqBAIJ*)A->data;
7   IS                iscol=a->col,isrow=a->row;
8   PetscErrorCode    ierr;
9   const PetscInt    *r,*c,*rout,*cout;
10   const PetscInt    *diag=a->diag,n=a->mbs,*vi,*ai=a->i,*aj=a->j;
11   PetscInt          i,nz,idx,idt,ii,ic,ir,oidx;
12   const MatScalar   *aa=a->a,*v;
13   PetscScalar       s1,s2,s3,s4,s5,s6,x1,x2,x3,x4,x5,x6,*x,*t;
14   const PetscScalar *b;
15 
16   PetscFunctionBegin;
17   ierr = VecGetArrayRead(bb,&b);CHKERRQ(ierr);
18   ierr = VecGetArray(xx,&x);CHKERRQ(ierr);
19   t    = a->solve_work;
20 
21   ierr = ISGetIndices(isrow,&rout);CHKERRQ(ierr); r = rout;
22   ierr = ISGetIndices(iscol,&cout);CHKERRQ(ierr); c = cout;
23 
24   /* copy the b into temp work space according to permutation */
25   ii = 0;
26   for (i=0; i<n; i++) {
27     ic      = 6*c[i];
28     t[ii]   = b[ic];
29     t[ii+1] = b[ic+1];
30     t[ii+2] = b[ic+2];
31     t[ii+3] = b[ic+3];
32     t[ii+4] = b[ic+4];
33     t[ii+5] = b[ic+5];
34     ii     += 6;
35   }
36 
37   /* forward solve the U^T */
38   idx = 0;
39   for (i=0; i<n; i++) {
40 
41     v = aa + 36*diag[i];
42     /* multiply by the inverse of the block diagonal */
43     x1 = t[idx];   x2 = t[1+idx]; x3    = t[2+idx]; x4 = t[3+idx]; x5 = t[4+idx];
44     x6 = t[5+idx];
45     s1 = v[0]*x1  +  v[1]*x2 +  v[2]*x3 +  v[3]*x4 +  v[4]*x5 +  v[5]*x6;
46     s2 = v[6]*x1  +  v[7]*x2 +  v[8]*x3 +  v[9]*x4 + v[10]*x5 + v[11]*x6;
47     s3 = v[12]*x1 + v[13]*x2 + v[14]*x3 + v[15]*x4 + v[16]*x5 + v[17]*x6;
48     s4 = v[18]*x1 + v[19]*x2 + v[20]*x3 + v[21]*x4 + v[22]*x5 + v[23]*x6;
49     s5 = v[24]*x1 + v[25]*x2 + v[26]*x3 + v[27]*x4 + v[28]*x5 + v[29]*x6;
50     s6 = v[30]*x1 + v[31]*x2 + v[32]*x3 + v[33]*x4 + v[34]*x5 + v[35]*x6;
51     v += 36;
52 
53     vi = aj + diag[i] + 1;
54     nz = ai[i+1] - diag[i] - 1;
55     while (nz--) {
56       oidx       = 6*(*vi++);
57       t[oidx]   -= v[0]*s1  +  v[1]*s2 +  v[2]*s3 +  v[3]*s4 +  v[4]*s5 +  v[5]*s6;
58       t[oidx+1] -= v[6]*s1  +  v[7]*s2 +  v[8]*s3 +  v[9]*s4 + v[10]*s5 + v[11]*s6;
59       t[oidx+2] -= v[12]*s1 + v[13]*s2 + v[14]*s3 + v[15]*s4 + v[16]*s5 + v[17]*s6;
60       t[oidx+3] -= v[18]*s1 + v[19]*s2 + v[20]*s3 + v[21]*s4 + v[22]*s5 + v[23]*s6;
61       t[oidx+4] -= v[24]*s1 + v[25]*s2 + v[26]*s3 + v[27]*s4 + v[28]*s5 + v[29]*s6;
62       t[oidx+5] -= v[30]*s1 + v[31]*s2 + v[32]*s3 + v[33]*s4 + v[34]*s5 + v[35]*s6;
63       v         += 36;
64     }
65     t[idx]   = s1;t[1+idx] = s2; t[2+idx] = s3;t[3+idx] = s4; t[4+idx] = s5;
66     t[5+idx] = s6;
67     idx     += 6;
68   }
69   /* backward solve the L^T */
70   for (i=n-1; i>=0; i--) {
71     v   = aa + 36*diag[i] - 36;
72     vi  = aj + diag[i] - 1;
73     nz  = diag[i] - ai[i];
74     idt = 6*i;
75     s1  = t[idt];  s2 = t[1+idt]; s3 = t[2+idt];s4 = t[3+idt]; s5 = t[4+idt];
76     s6  = t[5+idt];
77     while (nz--) {
78       idx       = 6*(*vi--);
79       t[idx]   -=  v[0]*s1 +  v[1]*s2 +  v[2]*s3 +  v[3]*s4 +  v[4]*s5 +  v[5]*s6;
80       t[idx+1] -=  v[6]*s1 +  v[7]*s2 +  v[8]*s3 +  v[9]*s4 + v[10]*s5 + v[11]*s6;
81       t[idx+2] -= v[12]*s1 + v[13]*s2 + v[14]*s3 + v[15]*s4 + v[16]*s5 + v[17]*s6;
82       t[idx+3] -= v[18]*s1 + v[19]*s2 + v[20]*s3 + v[21]*s4 + v[22]*s5 + v[23]*s6;
83       t[idx+4] -= v[24]*s1 + v[25]*s2 + v[26]*s3 + v[27]*s4 + v[28]*s5 + v[29]*s6;
84       t[idx+5] -= v[30]*s1 + v[31]*s2 + v[32]*s3 + v[33]*s4 + v[34]*s5 + v[35]*s6;
85       v        -= 36;
86     }
87   }
88 
89   /* copy t into x according to permutation */
90   ii = 0;
91   for (i=0; i<n; i++) {
92     ir      = 6*r[i];
93     x[ir]   = t[ii];
94     x[ir+1] = t[ii+1];
95     x[ir+2] = t[ii+2];
96     x[ir+3] = t[ii+3];
97     x[ir+4] = t[ii+4];
98     x[ir+5] = t[ii+5];
99     ii     += 6;
100   }
101 
102   ierr = ISRestoreIndices(isrow,&rout);CHKERRQ(ierr);
103   ierr = ISRestoreIndices(iscol,&cout);CHKERRQ(ierr);
104   ierr = VecRestoreArrayRead(bb,&b);CHKERRQ(ierr);
105   ierr = VecRestoreArray(xx,&x);CHKERRQ(ierr);
106   ierr = PetscLogFlops(2.0*36*(a->nz) - 6.0*A->cmap->n);CHKERRQ(ierr);
107   PetscFunctionReturn(0);
108 }
109 
110 PetscErrorCode MatSolveTranspose_SeqBAIJ_6(Mat A,Vec bb,Vec xx)
111 {
112   Mat_SeqBAIJ       *a=(Mat_SeqBAIJ*)A->data;
113   PetscErrorCode    ierr;
114   IS                iscol=a->col,isrow=a->row;
115   const PetscInt    n    =a->mbs,*vi,*ai=a->i,*aj=a->j,*diag=a->diag;
116   const PetscInt    *r,*c,*rout,*cout;
117   PetscInt          nz,idx,idt,j,i,oidx,ii,ic,ir;
118   const PetscInt    bs =A->rmap->bs,bs2=a->bs2;
119   const MatScalar   *aa=a->a,*v;
120   PetscScalar       s1,s2,s3,s4,s5,s6,x1,x2,x3,x4,x5,x6,*x,*t;
121   const PetscScalar *b;
122 
123   PetscFunctionBegin;
124   ierr = VecGetArrayRead(bb,&b);CHKERRQ(ierr);
125   ierr = VecGetArray(xx,&x);CHKERRQ(ierr);
126   t    = a->solve_work;
127 
128   ierr = ISGetIndices(isrow,&rout);CHKERRQ(ierr); r = rout;
129   ierr = ISGetIndices(iscol,&cout);CHKERRQ(ierr); c = cout;
130 
131   /* copy b into temp work space according to permutation */
132   for (i=0; i<n; i++) {
133     ii      = bs*i; ic = bs*c[i];
134     t[ii]   = b[ic]; t[ii+1] = b[ic+1]; t[ii+2] = b[ic+2]; t[ii+3] = b[ic+3];
135     t[ii+4] = b[ic+4];  t[ii+5] = b[ic+5];
136   }
137 
138   /* forward solve the U^T */
139   idx = 0;
140   for (i=0; i<n; i++) {
141     v = aa + bs2*diag[i];
142     /* multiply by the inverse of the block diagonal */
143     x1 = t[idx];   x2 = t[1+idx]; x3    = t[2+idx]; x4 = t[3+idx]; x5 = t[4+idx];
144     x6 = t[5+idx];
145     s1 = v[0]*x1  +  v[1]*x2 +  v[2]*x3 +  v[3]*x4 +  v[4]*x5 +  v[5]*x6;
146     s2 = v[6]*x1  +  v[7]*x2 +  v[8]*x3 +  v[9]*x4 + v[10]*x5 + v[11]*x6;
147     s3 = v[12]*x1 + v[13]*x2 + v[14]*x3 + v[15]*x4 + v[16]*x5 + v[17]*x6;
148     s4 = v[18]*x1 + v[19]*x2 + v[20]*x3 + v[21]*x4 + v[22]*x5 + v[23]*x6;
149     s5 = v[24]*x1 + v[25]*x2 + v[26]*x3 + v[27]*x4 + v[28]*x5 + v[29]*x6;
150     s6 = v[30]*x1 + v[31]*x2 + v[32]*x3 + v[33]*x4 + v[34]*x5 + v[35]*x6;
151     v -= bs2;
152 
153     vi = aj + diag[i] - 1;
154     nz = diag[i] - diag[i+1] - 1;
155     for (j=0; j>-nz; j--) {
156       oidx       = bs*vi[j];
157       t[oidx]   -= v[0]*s1  +  v[1]*s2 +  v[2]*s3 +  v[3]*s4 +  v[4]*s5 +  v[5]*s6;
158       t[oidx+1] -= v[6]*s1  +  v[7]*s2 +  v[8]*s3 +  v[9]*s4 + v[10]*s5 + v[11]*s6;
159       t[oidx+2] -= v[12]*s1 + v[13]*s2 + v[14]*s3 + v[15]*s4 + v[16]*s5 + v[17]*s6;
160       t[oidx+3] -= v[18]*s1 + v[19]*s2 + v[20]*s3 + v[21]*s4 + v[22]*s5 + v[23]*s6;
161       t[oidx+4] -= v[24]*s1 + v[25]*s2 + v[26]*s3 + v[27]*s4 + v[28]*s5 + v[29]*s6;
162       t[oidx+5] -= v[30]*s1 + v[31]*s2 + v[32]*s3 + v[33]*s4 + v[34]*s5 + v[35]*s6;
163       v         -= bs2;
164     }
165     t[idx]   = s1;t[1+idx] = s2;  t[2+idx] = s3;  t[3+idx] = s4; t[4+idx] =s5;
166     t[5+idx] = s6;
167     idx     += bs;
168   }
169   /* backward solve the L^T */
170   for (i=n-1; i>=0; i--) {
171     v   = aa + bs2*ai[i];
172     vi  = aj + ai[i];
173     nz  = ai[i+1] - ai[i];
174     idt = bs*i;
175     s1  = t[idt];  s2 = t[1+idt];  s3 = t[2+idt];  s4 = t[3+idt]; s5 = t[4+idt];
176     s6  = t[5+idt];
177     for (j=0; j<nz; j++) {
178       idx       = bs*vi[j];
179       t[idx]   -= v[0]*s1  +  v[1]*s2 +  v[2]*s3 +  v[3]*s4 +  v[4]*s5 +  v[5]*s6;
180       t[idx+1] -= v[6]*s1  +  v[7]*s2 +  v[8]*s3 +  v[9]*s4 + v[10]*s5 + v[11]*s6;
181       t[idx+2] -= v[12]*s1 + v[13]*s2 + v[14]*s3 + v[15]*s4 + v[16]*s5 + v[17]*s6;
182       t[idx+3] -= v[18]*s1 + v[19]*s2 + v[20]*s3 + v[21]*s4 + v[22]*s5 + v[23]*s6;
183       t[idx+4] -= v[24]*s1 + v[25]*s2 + v[26]*s3 + v[27]*s4 + v[28]*s5 + v[29]*s6;
184       t[idx+5] -= v[30]*s1 + v[31]*s2 + v[32]*s3 + v[33]*s4 + v[34]*s5 + v[35]*s6;
185       v        += bs2;
186     }
187   }
188 
189   /* copy t into x according to permutation */
190   for (i=0; i<n; i++) {
191     ii      = bs*i;  ir = bs*r[i];
192     x[ir]   = t[ii];  x[ir+1] = t[ii+1]; x[ir+2] = t[ii+2];  x[ir+3] = t[ii+3];
193     x[ir+4] = t[ii+4];  x[ir+5] = t[ii+5];
194   }
195 
196   ierr = ISRestoreIndices(isrow,&rout);CHKERRQ(ierr);
197   ierr = ISRestoreIndices(iscol,&cout);CHKERRQ(ierr);
198   ierr = VecRestoreArrayRead(bb,&b);CHKERRQ(ierr);
199   ierr = VecRestoreArray(xx,&x);CHKERRQ(ierr);
200   ierr = PetscLogFlops(2.0*bs2*(a->nz) - bs*A->cmap->n);CHKERRQ(ierr);
201   PetscFunctionReturn(0);
202 }
203