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