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