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