1 #include <petscdmda.h> 2 #include <petscts.h> 3 #include <adolc/interfaces.h> 4 #include "contexts.cxx" 5 6 /* 7 REQUIRES configuration of PETSc with option --download-adolc. 8 9 For documentation on ADOL-C, see 10 $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf 11 */ 12 13 /* 14 ADOL-C implementation for Jacobian vector product, using the forward mode of AD. 15 Intended to overload MatMult in matrix-free methods where implicit timestepping 16 has been used. 17 18 For an implicit Jacobian we may use the rule that 19 G = M*xdot + f(x) ==> dG/dx = a*M + df/dx, 20 where a = d(xdot)/dx is a constant. Evaluated at x0 and acting upon a vector x1: 21 (dG/dx)(x0) * x1 = (a*df/d(xdot)(x0) + df/dx)(x0) * x1. 22 23 Input parameters: 24 A_shell - Jacobian matrix of MatShell type 25 X - vector to be multiplied by A_shell 26 27 Output parameters: 28 Y - product of A_shell and X 29 */ 30 PetscErrorCode PetscAdolcIJacobianVectorProduct(Mat A_shell, Vec X, Vec Y) 31 { 32 AdolcMatCtx *mctx; 33 PetscInt m, n, i, j, k = 0, d; 34 const PetscScalar *x0; 35 PetscScalar *action, *x1; 36 Vec localX1; 37 DM da; 38 DMDALocalInfo info; 39 40 PetscFunctionBegin; 41 42 /* Get matrix-free context info */ 43 PetscCall(MatShellGetContext(A_shell, &mctx)); 44 m = mctx->m; 45 n = mctx->n; 46 47 /* Get local input vectors and extract data, x0 and x1*/ 48 PetscCall(TSGetDM(mctx->ts, &da)); 49 PetscCall(DMDAGetLocalInfo(da, &info)); 50 PetscCall(DMGetLocalVector(da, &localX1)); 51 PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX1)); 52 PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX1)); 53 54 PetscCall(VecGetArrayRead(mctx->localX0, &x0)); 55 PetscCall(VecGetArray(localX1, &x1)); 56 57 /* dF/dx part */ 58 PetscCall(PetscMalloc1(m, &action)); 59 PetscCall(PetscLogEventBegin(mctx->event1, 0, 0, 0, 0)); 60 fos_forward(mctx->tag1, m, n, 0, x0, x1, NULL, action); 61 for (j = info.gys; j < info.gys + info.gym; j++) { 62 for (i = info.gxs; i < info.gxs + info.gxm; i++) { 63 for (d = 0; d < 2; d++) { 64 if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(Y, 1, &k, &action[k], INSERT_VALUES)); 65 k++; 66 } 67 } 68 } 69 PetscCall(PetscLogEventEnd(mctx->event1, 0, 0, 0, 0)); 70 k = 0; 71 PetscCall(VecAssemblyBegin(Y)); /* Note: Need to assemble between separate calls */ 72 PetscCall(VecAssemblyEnd(Y)); /* to INSERT_VALUES and ADD_VALUES */ 73 74 /* a * dF/d(xdot) part */ 75 PetscCall(PetscLogEventBegin(mctx->event2, 0, 0, 0, 0)); 76 fos_forward(mctx->tag2, m, n, 0, x0, x1, NULL, action); 77 for (j = info.gys; j < info.gys + info.gym; j++) { 78 for (i = info.gxs; i < info.gxs + info.gxm; i++) { 79 for (d = 0; d < 2; d++) { 80 if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) { 81 action[k] *= mctx->shift; 82 PetscCall(VecSetValuesLocal(Y, 1, &k, &action[k], ADD_VALUES)); 83 } 84 k++; 85 } 86 } 87 } 88 PetscCall(PetscLogEventEnd(mctx->event2, 0, 0, 0, 0)); 89 PetscCall(VecAssemblyBegin(Y)); 90 PetscCall(VecAssemblyEnd(Y)); 91 PetscCall(PetscFree(action)); 92 93 /* Restore local vector */ 94 PetscCall(VecRestoreArray(localX1, &x1)); 95 PetscCall(VecRestoreArrayRead(mctx->localX0, &x0)); 96 PetscCall(DMRestoreLocalVector(da, &localX1)); 97 PetscFunctionReturn(PETSC_SUCCESS); 98 } 99 100 /* 101 ADOL-C implementation for Jacobian vector product, using the forward mode of AD. 102 Intended to overload MatMult in matrix-free methods where implicit timestepping 103 has been applied to a problem of the form 104 du/dt = F(u). 105 106 Input parameters: 107 A_shell - Jacobian matrix of MatShell type 108 X - vector to be multiplied by A_shell 109 110 Output parameters: 111 Y - product of A_shell and X 112 */ 113 PetscErrorCode PetscAdolcIJacobianVectorProductIDMass(Mat A_shell, Vec X, Vec Y) 114 { 115 AdolcMatCtx *mctx; 116 PetscInt m, n, i, j, k = 0, d; 117 const PetscScalar *x0; 118 PetscScalar *action, *x1; 119 Vec localX1; 120 DM da; 121 DMDALocalInfo info; 122 123 PetscFunctionBegin; 124 125 /* Get matrix-free context info */ 126 PetscCall(MatShellGetContext(A_shell, &mctx)); 127 m = mctx->m; 128 n = mctx->n; 129 130 /* Get local input vectors and extract data, x0 and x1*/ 131 PetscCall(TSGetDM(mctx->ts, &da)); 132 PetscCall(DMDAGetLocalInfo(da, &info)); 133 PetscCall(DMGetLocalVector(da, &localX1)); 134 PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX1)); 135 PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX1)); 136 137 PetscCall(VecGetArrayRead(mctx->localX0, &x0)); 138 PetscCall(VecGetArray(localX1, &x1)); 139 140 /* dF/dx part */ 141 PetscCall(PetscMalloc1(m, &action)); 142 PetscCall(PetscLogEventBegin(mctx->event1, 0, 0, 0, 0)); 143 fos_forward(mctx->tag1, m, n, 0, x0, x1, NULL, action); 144 for (j = info.gys; j < info.gys + info.gym; j++) { 145 for (i = info.gxs; i < info.gxs + info.gxm; i++) { 146 for (d = 0; d < 2; d++) { 147 if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(Y, 1, &k, &action[k], INSERT_VALUES)); 148 k++; 149 } 150 } 151 } 152 PetscCall(PetscLogEventEnd(mctx->event1, 0, 0, 0, 0)); 153 k = 0; 154 PetscCall(VecAssemblyBegin(Y)); /* Note: Need to assemble between separate calls */ 155 PetscCall(VecAssemblyEnd(Y)); /* to INSERT_VALUES and ADD_VALUES */ 156 PetscCall(PetscFree(action)); 157 158 /* Restore local vector */ 159 PetscCall(VecRestoreArray(localX1, &x1)); 160 PetscCall(VecRestoreArrayRead(mctx->localX0, &x0)); 161 PetscCall(DMRestoreLocalVector(da, &localX1)); 162 163 /* a * dF/d(xdot) part */ 164 PetscCall(PetscLogEventBegin(mctx->event2, 0, 0, 0, 0)); 165 PetscCall(VecAXPY(Y, mctx->shift, X)); 166 PetscCall(PetscLogEventEnd(mctx->event2, 0, 0, 0, 0)); 167 PetscFunctionReturn(PETSC_SUCCESS); 168 } 169 170 /* 171 ADOL-C implementation for Jacobian transpose vector product, using the reverse mode of AD. 172 Intended to overload MatMultTranspose in matrix-free methods where implicit timestepping 173 has been used. 174 175 Input parameters: 176 A_shell - Jacobian matrix of MatShell type 177 Y - vector to be multiplied by A_shell transpose 178 179 Output parameters: 180 X - product of A_shell transpose and X 181 */ 182 PetscErrorCode PetscAdolcIJacobianTransposeVectorProduct(Mat A_shell, Vec Y, Vec X) 183 { 184 AdolcMatCtx *mctx; 185 PetscInt m, n, i, j, k = 0, d; 186 const PetscScalar *x; 187 PetscScalar *action, *y; 188 Vec localY; 189 DM da; 190 DMDALocalInfo info; 191 192 PetscFunctionBegin; 193 194 /* Get matrix-free context info */ 195 PetscCall(MatShellGetContext(A_shell, &mctx)); 196 m = mctx->m; 197 n = mctx->n; 198 199 /* Get local input vectors and extract data, x0 and x1*/ 200 PetscCall(TSGetDM(mctx->ts, &da)); 201 PetscCall(DMDAGetLocalInfo(da, &info)); 202 PetscCall(DMGetLocalVector(da, &localY)); 203 PetscCall(DMGlobalToLocalBegin(da, Y, INSERT_VALUES, localY)); 204 PetscCall(DMGlobalToLocalEnd(da, Y, INSERT_VALUES, localY)); 205 PetscCall(VecGetArrayRead(mctx->localX0, &x)); 206 PetscCall(VecGetArray(localY, &y)); 207 208 /* dF/dx part */ 209 PetscCall(PetscMalloc1(n, &action)); 210 PetscCall(PetscLogEventBegin(mctx->event3, 0, 0, 0, 0)); 211 if (!mctx->flg) zos_forward(mctx->tag1, m, n, 1, x, NULL); 212 fos_reverse(mctx->tag1, m, n, y, action); 213 for (j = info.gys; j < info.gys + info.gym; j++) { 214 for (i = info.gxs; i < info.gxs + info.gxm; i++) { 215 for (d = 0; d < 2; d++) { 216 if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(X, 1, &k, &action[k], INSERT_VALUES)); 217 k++; 218 } 219 } 220 } 221 PetscCall(PetscLogEventEnd(mctx->event3, 0, 0, 0, 0)); 222 k = 0; 223 PetscCall(VecAssemblyBegin(X)); /* Note: Need to assemble between separate calls */ 224 PetscCall(VecAssemblyEnd(X)); /* to INSERT_VALUES and ADD_VALUES */ 225 226 /* a * dF/d(xdot) part */ 227 PetscCall(PetscLogEventBegin(mctx->event4, 0, 0, 0, 0)); 228 if (!mctx->flg) { 229 zos_forward(mctx->tag2, m, n, 1, x, NULL); 230 mctx->flg = PETSC_TRUE; 231 } 232 fos_reverse(mctx->tag2, m, n, y, action); 233 for (j = info.gys; j < info.gys + info.gym; j++) { 234 for (i = info.gxs; i < info.gxs + info.gxm; i++) { 235 for (d = 0; d < 2; d++) { 236 if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) { 237 action[k] *= mctx->shift; 238 PetscCall(VecSetValuesLocal(X, 1, &k, &action[k], ADD_VALUES)); 239 } 240 k++; 241 } 242 } 243 } 244 PetscCall(PetscLogEventEnd(mctx->event4, 0, 0, 0, 0)); 245 PetscCall(VecAssemblyBegin(X)); 246 PetscCall(VecAssemblyEnd(X)); 247 PetscCall(PetscFree(action)); 248 249 /* Restore local vector */ 250 PetscCall(VecRestoreArray(localY, &y)); 251 PetscCall(VecRestoreArrayRead(mctx->localX0, &x)); 252 PetscCall(DMRestoreLocalVector(da, &localY)); 253 PetscFunctionReturn(PETSC_SUCCESS); 254 } 255 256 /* 257 ADOL-C implementation for Jacobian transpose vector product, using the reverse mode of AD. 258 Intended to overload MatMultTranspose in matrix-free methods where implicit timestepping 259 has been applied to a problem of the form 260 du/dt = F(u). 261 262 Input parameters: 263 A_shell - Jacobian matrix of MatShell type 264 Y - vector to be multiplied by A_shell transpose 265 266 Output parameters: 267 X - product of A_shell transpose and X 268 */ 269 PetscErrorCode PetscAdolcIJacobianTransposeVectorProductIDMass(Mat A_shell, Vec Y, Vec X) 270 { 271 AdolcMatCtx *mctx; 272 PetscInt m, n, i, j, k = 0, d; 273 const PetscScalar *x; 274 PetscScalar *action, *y; 275 Vec localY; 276 DM da; 277 DMDALocalInfo info; 278 279 PetscFunctionBegin; 280 281 /* Get matrix-free context info */ 282 PetscCall(MatShellGetContext(A_shell, &mctx)); 283 m = mctx->m; 284 n = mctx->n; 285 286 /* Get local input vectors and extract data, x0 and x1*/ 287 PetscCall(TSGetDM(mctx->ts, &da)); 288 PetscCall(DMDAGetLocalInfo(da, &info)); 289 PetscCall(DMGetLocalVector(da, &localY)); 290 PetscCall(DMGlobalToLocalBegin(da, Y, INSERT_VALUES, localY)); 291 PetscCall(DMGlobalToLocalEnd(da, Y, INSERT_VALUES, localY)); 292 PetscCall(VecGetArrayRead(mctx->localX0, &x)); 293 PetscCall(VecGetArray(localY, &y)); 294 295 /* dF/dx part */ 296 PetscCall(PetscMalloc1(n, &action)); 297 PetscCall(PetscLogEventBegin(mctx->event3, 0, 0, 0, 0)); 298 if (!mctx->flg) zos_forward(mctx->tag1, m, n, 1, x, NULL); 299 fos_reverse(mctx->tag1, m, n, y, action); 300 for (j = info.gys; j < info.gys + info.gym; j++) { 301 for (i = info.gxs; i < info.gxs + info.gxm; i++) { 302 for (d = 0; d < 2; d++) { 303 if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(X, 1, &k, &action[k], INSERT_VALUES)); 304 k++; 305 } 306 } 307 } 308 PetscCall(PetscLogEventEnd(mctx->event3, 0, 0, 0, 0)); 309 k = 0; 310 PetscCall(VecAssemblyBegin(X)); /* Note: Need to assemble between separate calls */ 311 PetscCall(VecAssemblyEnd(X)); /* to INSERT_VALUES and ADD_VALUES */ 312 PetscCall(PetscFree(action)); 313 314 /* Restore local vector */ 315 PetscCall(VecRestoreArray(localY, &y)); 316 PetscCall(VecRestoreArrayRead(mctx->localX0, &x)); 317 PetscCall(DMRestoreLocalVector(da, &localY)); 318 319 /* a * dF/d(xdot) part */ 320 PetscCall(PetscLogEventBegin(mctx->event4, 0, 0, 0, 0)); 321 PetscCall(VecAXPY(X, mctx->shift, Y)); 322 PetscCall(PetscLogEventEnd(mctx->event4, 0, 0, 0, 0)); 323 PetscFunctionReturn(PETSC_SUCCESS); 324 } 325