1 #include "../include/matops.h" 2 #include "../include/petscutils.h" 3 4 // ----------------------------------------------------------------------------- 5 // Setup apply operator context data 6 // ----------------------------------------------------------------------------- 7 PetscErrorCode SetupApplyOperatorCtx(MPI_Comm comm, DM dm, Ceed ceed, 8 CeedData ceed_data, Vec X_loc, 9 OperatorApplyContext op_apply_ctx) { 10 PetscErrorCode ierr; 11 PetscFunctionBeginUser; 12 13 op_apply_ctx->comm = comm; 14 op_apply_ctx->dm = dm; 15 op_apply_ctx->X_loc = X_loc; 16 ierr = VecDuplicate(X_loc, &op_apply_ctx->Y_loc); CHKERRQ(ierr); 17 op_apply_ctx->x_ceed = ceed_data->x_ceed; 18 op_apply_ctx->y_ceed = ceed_data->y_ceed; 19 op_apply_ctx->op = ceed_data->op_apply; 20 op_apply_ctx->ceed = ceed; 21 PetscFunctionReturn(0); 22 } 23 24 // ----------------------------------------------------------------------------- 25 // Setup error operator context data 26 // ----------------------------------------------------------------------------- 27 PetscErrorCode SetupErrorOperatorCtx(MPI_Comm comm, DM dm, Ceed ceed, 28 CeedData ceed_data, Vec X_loc, CeedOperator op_error, 29 OperatorApplyContext op_error_ctx) { 30 PetscErrorCode ierr; 31 PetscFunctionBeginUser; 32 33 op_error_ctx->comm = comm; 34 op_error_ctx->dm = dm; 35 op_error_ctx->X_loc = X_loc; 36 ierr = VecDuplicate(X_loc, &op_error_ctx->Y_loc); CHKERRQ(ierr); 37 op_error_ctx->x_ceed = ceed_data->x_ceed; 38 op_error_ctx->y_ceed = ceed_data->y_ceed; 39 op_error_ctx->op = op_error; 40 op_error_ctx->ceed = ceed; 41 PetscFunctionReturn(0); 42 } 43 44 // ----------------------------------------------------------------------------- 45 // This function returns the computed diagonal of the operator 46 // ----------------------------------------------------------------------------- 47 PetscErrorCode MatGetDiag(Mat A, Vec D) { 48 PetscErrorCode ierr; 49 OperatorApplyContext op_apply_ctx; 50 51 PetscFunctionBeginUser; 52 53 ierr = MatShellGetContext(A, &op_apply_ctx); CHKERRQ(ierr); 54 55 // Compute Diagonal via libCEED 56 PetscScalar *y; 57 PetscMemType mem_type; 58 59 // -- Place PETSc vector in libCEED vector 60 ierr = VecGetArrayAndMemType(op_apply_ctx->Y_loc, &y, &mem_type); CHKERRQ(ierr); 61 CeedVectorSetArray(op_apply_ctx->y_ceed, MemTypeP2C(mem_type), 62 CEED_USE_POINTER, y); 63 64 // -- Compute Diagonal 65 CeedOperatorLinearAssembleDiagonal(op_apply_ctx->op, op_apply_ctx->y_ceed, 66 CEED_REQUEST_IMMEDIATE); 67 68 // -- Local-to-Global 69 CeedVectorTakeArray(op_apply_ctx->y_ceed, MemTypeP2C(mem_type), NULL); 70 ierr = VecRestoreArrayAndMemType(op_apply_ctx->Y_loc, &y); CHKERRQ(ierr); 71 ierr = VecZeroEntries(D); CHKERRQ(ierr); 72 ierr = DMLocalToGlobal(op_apply_ctx->dm, op_apply_ctx->Y_loc, ADD_VALUES, D); 73 CHKERRQ(ierr); 74 75 PetscFunctionReturn(0); 76 }; 77 78 // ----------------------------------------------------------------------------- 79 // This function uses libCEED to compute the action of the Laplacian with 80 // Dirichlet boundary conditions 81 // ----------------------------------------------------------------------------- 82 PetscErrorCode ApplyLocal_Ceed(Vec X, Vec Y, 83 OperatorApplyContext op_apply_ctx) { 84 PetscErrorCode ierr; 85 PetscScalar *x, *y; 86 PetscMemType x_mem_type, y_mem_type; 87 88 PetscFunctionBeginUser; 89 90 // Global-to-local 91 ierr = DMGlobalToLocal(op_apply_ctx->dm, X, INSERT_VALUES, op_apply_ctx->X_loc); 92 CHKERRQ(ierr); 93 94 // Setup libCEED vectors 95 ierr = VecGetArrayReadAndMemType(op_apply_ctx->X_loc, (const PetscScalar **)&x, 96 &x_mem_type); CHKERRQ(ierr); 97 ierr = VecGetArrayAndMemType(op_apply_ctx->Y_loc, &y, &y_mem_type); 98 CHKERRQ(ierr); 99 CeedVectorSetArray(op_apply_ctx->x_ceed, MemTypeP2C(x_mem_type), 100 CEED_USE_POINTER, x); 101 CeedVectorSetArray(op_apply_ctx->y_ceed, MemTypeP2C(y_mem_type), 102 CEED_USE_POINTER, y); 103 104 // Apply libCEED operator 105 CeedOperatorApply(op_apply_ctx->op, op_apply_ctx->x_ceed, op_apply_ctx->y_ceed, 106 CEED_REQUEST_IMMEDIATE); 107 108 // Restore PETSc vectors 109 CeedVectorTakeArray(op_apply_ctx->x_ceed, MemTypeP2C(x_mem_type), NULL); 110 CeedVectorTakeArray(op_apply_ctx->y_ceed, MemTypeP2C(y_mem_type), NULL); 111 ierr = VecRestoreArrayReadAndMemType(op_apply_ctx->X_loc, 112 (const PetscScalar **)&x); 113 CHKERRQ(ierr); 114 ierr = VecRestoreArrayAndMemType(op_apply_ctx->Y_loc, &y); CHKERRQ(ierr); 115 116 // Local-to-global 117 ierr = VecZeroEntries(Y); CHKERRQ(ierr); 118 ierr = DMLocalToGlobal(op_apply_ctx->dm, op_apply_ctx->Y_loc, ADD_VALUES, Y); 119 CHKERRQ(ierr); 120 121 PetscFunctionReturn(0); 122 }; 123 124 // ----------------------------------------------------------------------------- 125 // This function wraps the libCEED operator for a MatShell 126 // ----------------------------------------------------------------------------- 127 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { 128 PetscErrorCode ierr; 129 OperatorApplyContext op_apply_ctx; 130 131 PetscFunctionBeginUser; 132 133 ierr = MatShellGetContext(A, &op_apply_ctx); CHKERRQ(ierr); 134 135 // libCEED for local action of residual evaluator 136 ierr = ApplyLocal_Ceed(X, Y, op_apply_ctx); CHKERRQ(ierr); 137 138 PetscFunctionReturn(0); 139 }; 140 141 // ----------------------------------------------------------------------------- 142 // This function wraps the libCEED operator for a SNES residual evaluation 143 // ----------------------------------------------------------------------------- 144 PetscErrorCode FormResidual_Ceed(SNES snes, Vec X, Vec Y, void *ctx) { 145 PetscErrorCode ierr; 146 OperatorApplyContext op_apply_ctx = (OperatorApplyContext)ctx; 147 148 PetscFunctionBeginUser; 149 150 // libCEED for local action of residual evaluator 151 ierr = ApplyLocal_Ceed(X, Y, op_apply_ctx); CHKERRQ(ierr); 152 153 PetscFunctionReturn(0); 154 }; 155 156 // ----------------------------------------------------------------------------- 157 // This function uses libCEED to compute the action of the prolongation operator 158 // ----------------------------------------------------------------------------- 159 PetscErrorCode MatMult_Prolong(Mat A, Vec X, Vec Y) { 160 PetscErrorCode ierr; 161 ProlongRestrContext pr_restr_ctx; 162 PetscScalar *c, *f; 163 PetscMemType c_mem_type, f_mem_type; 164 165 PetscFunctionBeginUser; 166 167 ierr = MatShellGetContext(A, &pr_restr_ctx); CHKERRQ(ierr); 168 169 // Global-to-local 170 ierr = VecZeroEntries(pr_restr_ctx->loc_vec_c); CHKERRQ(ierr); 171 ierr = DMGlobalToLocal(pr_restr_ctx->dmc, X, INSERT_VALUES, 172 pr_restr_ctx->loc_vec_c); 173 CHKERRQ(ierr); 174 175 // Setup libCEED vectors 176 ierr = VecGetArrayReadAndMemType(pr_restr_ctx->loc_vec_c, 177 (const PetscScalar **)&c, 178 &c_mem_type); CHKERRQ(ierr); 179 ierr = VecGetArrayAndMemType(pr_restr_ctx->loc_vec_f, &f, &f_mem_type); 180 CHKERRQ(ierr); 181 CeedVectorSetArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), 182 CEED_USE_POINTER, c); 183 CeedVectorSetArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), 184 CEED_USE_POINTER, f); 185 186 // Apply libCEED operator 187 CeedOperatorApply(pr_restr_ctx->op_prolong, pr_restr_ctx->ceed_vec_c, 188 pr_restr_ctx->ceed_vec_f, CEED_REQUEST_IMMEDIATE); 189 190 // Restore PETSc vectors 191 CeedVectorTakeArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), NULL); 192 CeedVectorTakeArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), NULL); 193 ierr = VecRestoreArrayReadAndMemType(pr_restr_ctx->loc_vec_c, 194 (const PetscScalar **)&c); 195 CHKERRQ(ierr); 196 ierr = VecRestoreArrayAndMemType(pr_restr_ctx->loc_vec_f, &f); CHKERRQ(ierr); 197 198 // Multiplicity 199 ierr = VecPointwiseMult(pr_restr_ctx->loc_vec_f, pr_restr_ctx->loc_vec_f, 200 pr_restr_ctx->mult_vec); 201 202 // Local-to-global 203 ierr = VecZeroEntries(Y); CHKERRQ(ierr); 204 ierr = DMLocalToGlobal(pr_restr_ctx->dmf, pr_restr_ctx->loc_vec_f, ADD_VALUES, 205 Y); CHKERRQ(ierr); 206 207 PetscFunctionReturn(0); 208 }; 209 210 // ----------------------------------------------------------------------------- 211 // This function uses libCEED to compute the action of the restriction operator 212 // ----------------------------------------------------------------------------- 213 PetscErrorCode MatMult_Restrict(Mat A, Vec X, Vec Y) { 214 PetscErrorCode ierr; 215 ProlongRestrContext pr_restr_ctx; 216 PetscScalar *c, *f; 217 PetscMemType c_mem_type, f_mem_type; 218 219 PetscFunctionBeginUser; 220 221 ierr = MatShellGetContext(A, &pr_restr_ctx); CHKERRQ(ierr); 222 223 // Global-to-local 224 ierr = VecZeroEntries(pr_restr_ctx->loc_vec_f); CHKERRQ(ierr); 225 ierr = DMGlobalToLocal(pr_restr_ctx->dmf, X, INSERT_VALUES, 226 pr_restr_ctx->loc_vec_f); 227 CHKERRQ(ierr); 228 229 // Multiplicity 230 ierr = VecPointwiseMult(pr_restr_ctx->loc_vec_f, pr_restr_ctx->loc_vec_f, 231 pr_restr_ctx->mult_vec); 232 CHKERRQ(ierr); 233 234 // Setup libCEED vectors 235 ierr = VecGetArrayReadAndMemType(pr_restr_ctx->loc_vec_f, 236 (const PetscScalar **)&f, 237 &f_mem_type); CHKERRQ(ierr); 238 ierr = VecGetArrayAndMemType(pr_restr_ctx->loc_vec_c, &c, &c_mem_type); 239 CHKERRQ(ierr); 240 CeedVectorSetArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), 241 CEED_USE_POINTER, f); 242 CeedVectorSetArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), 243 CEED_USE_POINTER, c); 244 245 // Apply CEED operator 246 CeedOperatorApply(pr_restr_ctx->op_restrict, pr_restr_ctx->ceed_vec_f, 247 pr_restr_ctx->ceed_vec_c, CEED_REQUEST_IMMEDIATE); 248 249 // Restore PETSc vectors 250 CeedVectorTakeArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), NULL); 251 CeedVectorTakeArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), NULL); 252 ierr = VecRestoreArrayReadAndMemType(pr_restr_ctx->loc_vec_f, 253 (const PetscScalar **)&f); CHKERRQ(ierr); 254 ierr = VecRestoreArrayAndMemType(pr_restr_ctx->loc_vec_c, &c); CHKERRQ(ierr); 255 256 // Local-to-global 257 ierr = VecZeroEntries(Y); CHKERRQ(ierr); 258 ierr = DMLocalToGlobal(pr_restr_ctx->dmc, pr_restr_ctx->loc_vec_c, ADD_VALUES, 259 Y); 260 CHKERRQ(ierr); 261 262 PetscFunctionReturn(0); 263 }; 264 265 // ----------------------------------------------------------------------------- 266 // This function calculates the error in the final solution 267 // ----------------------------------------------------------------------------- 268 PetscErrorCode ComputeErrorMax(OperatorApplyContext op_error_ctx, 269 Vec X, CeedVector target, 270 PetscScalar *max_error) { 271 PetscErrorCode ierr; 272 PetscScalar *x; 273 PetscMemType mem_type; 274 CeedVector collocated_error; 275 CeedSize length; 276 277 PetscFunctionBeginUser; 278 CeedVectorGetLength(target, &length); 279 CeedVectorCreate(op_error_ctx->ceed, length, &collocated_error); 280 281 // Global-to-local 282 ierr = DMGlobalToLocal(op_error_ctx->dm, X, INSERT_VALUES, op_error_ctx->X_loc); 283 CHKERRQ(ierr); 284 285 // Setup CEED vector 286 ierr = VecGetArrayAndMemType(op_error_ctx->X_loc, &x, &mem_type); CHKERRQ(ierr); 287 CeedVectorSetArray(op_error_ctx->x_ceed, MemTypeP2C(mem_type), 288 CEED_USE_POINTER, x); 289 290 // Apply CEED operator 291 CeedOperatorApply(op_error_ctx->op, op_error_ctx->x_ceed, collocated_error, 292 CEED_REQUEST_IMMEDIATE); 293 294 // Restore PETSc vector 295 CeedVectorTakeArray(op_error_ctx->x_ceed, MemTypeP2C(mem_type), NULL); 296 ierr = VecRestoreArrayReadAndMemType(op_error_ctx->X_loc, 297 (const PetscScalar **)&x); 298 CHKERRQ(ierr); 299 300 // Reduce max error 301 *max_error = 0; 302 const CeedScalar *e; 303 CeedVectorGetArrayRead(collocated_error, CEED_MEM_HOST, &e); 304 for (CeedInt i=0; i<length; i++) { 305 *max_error = PetscMax(*max_error, PetscAbsScalar(e[i])); 306 } 307 CeedVectorRestoreArrayRead(collocated_error, &e); 308 ierr = MPI_Allreduce(MPI_IN_PLACE, max_error, 1, MPIU_REAL, MPIU_MAX, 309 op_error_ctx->comm); 310 CHKERRQ(ierr); 311 312 // Cleanup 313 CeedVectorDestroy(&collocated_error); 314 315 PetscFunctionReturn(0); 316 }; 317 318 // ----------------------------------------------------------------------------- 319