1 #include <petscconf.h> 2 3 #if defined(PETSC_HAVE_KOKKOS_KERNELS) 4 #include <Kokkos_Core.hpp> 5 #include <petscdmda_kokkos.hpp> 6 7 #include <petscdm.h> 8 #include <petscdmda.h> 9 #include <petscsnes.h> 10 #include "ex55.h" 11 12 using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space; 13 using ConstPetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<const PetscScalar **, Kokkos::LayoutRight, DefaultMemorySpace>; 14 using PetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<PetscScalar **, Kokkos::LayoutRight, DefaultMemorySpace>; 15 16 using PetscCountKokkosView = Kokkos::View<PetscCount *, DefaultMemorySpace>; 17 using PetscIntKokkosView = Kokkos::View<PetscInt *, DefaultMemorySpace>; 18 using PetscScalarKokkosView = Kokkos::View<PetscScalar *, DefaultMemorySpace>; 19 using Kokkos::Iterate; 20 using Kokkos::MDRangePolicy; 21 using Kokkos::Rank; 22 23 KOKKOS_INLINE_FUNCTION PetscErrorCode MMSSolution1(AppCtx *user, const DMDACoor2d *c, PetscScalar *u) 24 { 25 PetscReal x = PetscRealPart(c->x), y = PetscRealPart(c->y); 26 u[0] = x * (1 - x) * y * (1 - y); 27 return PETSC_SUCCESS; 28 } 29 30 KOKKOS_INLINE_FUNCTION PetscErrorCode MMSForcing1(PetscReal user_param, const DMDACoor2d *c, PetscScalar *f) 31 { 32 PetscReal x = PetscRealPart(c->x), y = PetscRealPart(c->y); 33 f[0] = 2 * x * (1 - x) + 2 * y * (1 - y) - user_param * PetscExpReal(x * (1 - x) * y * (1 - y)); 34 return PETSC_SUCCESS; 35 } 36 37 PetscErrorCode FormFunctionLocalVec(DMDALocalInfo *info, Vec x, Vec f, AppCtx *user) 38 { 39 PetscReal lambda, hx, hy, hxdhy, hydhx; 40 PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my; 41 PetscReal user_param = user->param; 42 43 ConstPetscScalarKokkosOffsetView2D xv; 44 PetscScalarKokkosOffsetView2D fv; 45 46 PetscFunctionBeginUser; 47 lambda = user->param; 48 hx = 1.0 / (PetscReal)(info->mx - 1); 49 hy = 1.0 / (PetscReal)(info->my - 1); 50 hxdhy = hx / hy; 51 hydhx = hy / hx; 52 /* 53 Compute function over the locally owned part of the grid 54 */ 55 PetscCall(DMDAVecGetKokkosOffsetView(info->da, x, &xv)); 56 PetscCall(DMDAVecGetKokkosOffsetViewWrite(info->da, f, &fv)); 57 58 PetscCallCXX(Kokkos::parallel_for( 59 "FormFunctionLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscInt j, PetscInt i) { 60 DMDACoor2d c; 61 PetscScalar u, ue, uw, un, us, uxx, uyy, mms_solution, mms_forcing; 62 63 if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) { 64 c.x = i * hx; 65 c.y = j * hy; 66 static_cast<void>(MMSSolution1(user, &c, &mms_solution)); 67 fv(j, i) = 2.0 * (hydhx + hxdhy) * (xv(j, i) - mms_solution); 68 } else { 69 u = xv(j, i); 70 uw = xv(j, i - 1); 71 ue = xv(j, i + 1); 72 un = xv(j - 1, i); 73 us = xv(j + 1, i); 74 75 /* Enforce boundary conditions at neighboring points -- setting these values causes the Jacobian to be symmetric. */ 76 if (i - 1 == 0) { 77 c.x = (i - 1) * hx; 78 c.y = j * hy; 79 static_cast<void>(MMSSolution1(user, &c, &uw)); 80 } 81 if (i + 1 == mx - 1) { 82 c.x = (i + 1) * hx; 83 c.y = j * hy; 84 static_cast<void>(MMSSolution1(user, &c, &ue)); 85 } 86 if (j - 1 == 0) { 87 c.x = i * hx; 88 c.y = (j - 1) * hy; 89 static_cast<void>(MMSSolution1(user, &c, &un)); 90 } 91 if (j + 1 == my - 1) { 92 c.x = i * hx; 93 c.y = (j + 1) * hy; 94 static_cast<void>(MMSSolution1(user, &c, &us)); 95 } 96 97 uxx = (2.0 * u - uw - ue) * hydhx; 98 uyy = (2.0 * u - un - us) * hxdhy; 99 mms_forcing = 0; 100 c.x = i * hx; 101 c.y = j * hy; 102 static_cast<void>(MMSForcing1(user_param, &c, &mms_forcing)); 103 fv(j, i) = uxx + uyy - hx * hy * (lambda * PetscExpScalar(u) + mms_forcing); 104 } 105 })); 106 107 PetscCall(DMDAVecRestoreKokkosOffsetView(info->da, x, &xv)); 108 PetscCall(DMDAVecRestoreKokkosOffsetViewWrite(info->da, f, &fv)); 109 110 PetscCall(PetscLogFlops(11.0 * info->ym * info->xm)); 111 PetscFunctionReturn(PETSC_SUCCESS); 112 } 113 114 PetscErrorCode FormObjectiveLocalVec(DMDALocalInfo *info, Vec x, PetscReal *obj, AppCtx *user) 115 { 116 PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my; 117 PetscReal lambda, hx, hy, hxdhy, hydhx, sc, lobj = 0; 118 MPI_Comm comm; 119 120 ConstPetscScalarKokkosOffsetView2D xv; 121 122 PetscFunctionBeginUser; 123 *obj = 0; 124 PetscCall(PetscObjectGetComm((PetscObject)info->da, &comm)); 125 lambda = user->param; 126 hx = 1.0 / (PetscReal)(mx - 1); 127 hy = 1.0 / (PetscReal)(my - 1); 128 sc = hx * hy * lambda; 129 hxdhy = hx / hy; 130 hydhx = hy / hx; 131 /* 132 Compute function over the locally owned part of the grid 133 */ 134 PetscCall(DMDAVecGetKokkosOffsetView(info->da, x, &xv)); 135 136 PetscCallCXX(Kokkos::parallel_reduce( 137 "FormObjectiveLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), 138 KOKKOS_LAMBDA(PetscInt j, PetscInt i, PetscReal &update) { 139 PetscScalar u, ue, uw, un, us, uxux, uyuy; 140 if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) { 141 update += PetscRealPart((hydhx + hxdhy) * xv(j, i) * xv(j, i)); 142 } else { 143 u = xv(j, i); 144 uw = xv(j, i - 1); 145 ue = xv(j, i + 1); 146 un = xv(j - 1, i); 147 us = xv(j + 1, i); 148 149 if (i - 1 == 0) uw = 0.; 150 if (i + 1 == mx - 1) ue = 0.; 151 if (j - 1 == 0) un = 0.; 152 if (j + 1 == my - 1) us = 0.; 153 154 /* F[u] = 1/2\int_{\omega}\nabla^2u(x)*u(x)*dx */ 155 156 uxux = u * (2. * u - ue - uw) * hydhx; 157 uyuy = u * (2. * u - un - us) * hxdhy; 158 159 update += PetscRealPart(0.5 * (uxux + uyuy) - sc * PetscExpScalar(u)); 160 } 161 }, 162 lobj)); 163 164 PetscCall(DMDAVecRestoreKokkosOffsetView(info->da, x, &xv)); 165 PetscCall(PetscLogFlops(12.0 * info->ym * info->xm)); 166 *obj = lobj; 167 PetscFunctionReturn(PETSC_SUCCESS); 168 } 169 170 PetscErrorCode FormJacobianLocalVec(DMDALocalInfo *info, Vec x, Mat jac, Mat jacpre, AppCtx *user) 171 { 172 PetscInt i, j; 173 PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my; 174 MatStencil col[5], row; 175 PetscScalar lambda, hx, hy, hxdhy, hydhx, sc; 176 DM coordDA; 177 Vec coordinates; 178 DMDACoor2d **coords; 179 180 PetscFunctionBeginUser; 181 lambda = user->param; 182 /* Extract coordinates */ 183 PetscCall(DMGetCoordinateDM(info->da, &coordDA)); 184 PetscCall(DMGetCoordinates(info->da, &coordinates)); 185 186 PetscCall(DMDAVecGetArray(coordDA, coordinates, &coords)); 187 hx = xm > 1 ? PetscRealPart(coords[ys][xs + 1].x) - PetscRealPart(coords[ys][xs].x) : 1.0; 188 hy = ym > 1 ? PetscRealPart(coords[ys + 1][xs].y) - PetscRealPart(coords[ys][xs].y) : 1.0; 189 PetscCall(DMDAVecRestoreArray(coordDA, coordinates, &coords)); 190 191 hxdhy = hx / hy; 192 hydhx = hy / hx; 193 sc = hx * hy * lambda; 194 195 /* ----------------------------------------- */ 196 /* MatSetPreallocationCOO() */ 197 /* ----------------------------------------- */ 198 if (!user->ncoo) { 199 PetscCount ncoo = ((PetscCount)xm) * ((PetscCount)ym) * 5; 200 PetscInt *coo_i, *coo_j, *ip, *jp; 201 PetscCall(PetscMalloc2(ncoo, &coo_i, ncoo, &coo_j)); /* 5-point stencil such that each row has at most 5 nonzeros */ 202 203 ip = coo_i; 204 jp = coo_j; 205 for (j = ys; j < ys + ym; j++) { 206 for (i = xs; i < xs + xm; i++) { 207 row.j = j; 208 row.i = i; 209 /* Initialize neighbors with negative indices */ 210 col[0].j = col[1].j = col[3].j = col[4].j = -1; 211 /* boundary points */ 212 if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) { 213 col[2].j = row.j; 214 col[2].i = row.i; 215 } else { 216 /* interior grid points */ 217 if (j - 1 != 0) { 218 col[0].j = j - 1; 219 col[0].i = i; 220 } 221 222 if (i - 1 != 0) { 223 col[1].j = j; 224 col[1].i = i - 1; 225 } 226 227 col[2].j = row.j; 228 col[2].i = row.i; 229 230 if (i + 1 != mx - 1) { 231 col[3].j = j; 232 col[3].i = i + 1; 233 } 234 235 if (j + 1 != mx - 1) { 236 col[4].j = j + 1; 237 col[4].i = i; 238 } 239 } 240 PetscCall(DMDAMapMatStencilToGlobal(info->da, 5, col, jp)); 241 for (PetscInt k = 0; k < 5; k++) ip[k] = jp[2]; 242 ip += 5; 243 jp += 5; 244 } 245 } 246 PetscCall(MatSetPreallocationCOO(jacpre, ncoo, coo_i, coo_j)); 247 PetscCall(PetscFree2(coo_i, coo_j)); 248 user->ncoo = ncoo; 249 } 250 251 /* ----------------------------------------- */ 252 /* MatSetValuesCOO() */ 253 /* ----------------------------------------- */ 254 PetscScalarKokkosView coo_v("coo_v", user->ncoo); 255 ConstPetscScalarKokkosOffsetView2D xv; 256 257 PetscCall(DMDAVecGetKokkosOffsetView(info->da, x, &xv)); 258 259 PetscCallCXX(Kokkos::parallel_for( 260 "FormFunctionLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscCount j, PetscCount i) { 261 PetscInt p = ((j - ys) * xm + (i - xs)) * 5; 262 /* boundary points */ 263 if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) { 264 coo_v(p + 2) = 2.0 * (hydhx + hxdhy); 265 } else { 266 /* interior grid points */ 267 if (j - 1 != 0) coo_v(p + 0) = -hxdhy; 268 if (i - 1 != 0) coo_v(p + 1) = -hydhx; 269 270 coo_v(p + 2) = 2.0 * (hydhx + hxdhy) - sc * PetscExpScalar(xv(j, i)); 271 272 if (i + 1 != mx - 1) coo_v(p + 3) = -hydhx; 273 if (j + 1 != mx - 1) coo_v(p + 4) = -hxdhy; 274 } 275 })); 276 PetscCall(MatSetValuesCOO(jacpre, coo_v.data(), INSERT_VALUES)); 277 PetscCall(DMDAVecRestoreKokkosOffsetView(info->da, x, &xv)); 278 PetscFunctionReturn(PETSC_SUCCESS); 279 } 280 281 #else 282 #include "ex55.h" 283 284 PetscErrorCode FormObjectiveLocalVec(DMDALocalInfo *info, Vec x, PetscReal *obj, AppCtx *user) 285 { 286 PetscFunctionBeginUser; 287 PetscCheck(PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_SUP, "Need to reconfigure with --download-kokkos-kernels"); 288 PetscFunctionReturn(PETSC_SUCCESS); 289 } 290 291 PetscErrorCode FormFunctionLocalVec(DMDALocalInfo *info, Vec x, Vec f, AppCtx *user) 292 { 293 PetscFunctionBeginUser; 294 PetscCheck(PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_SUP, "Need to reconfigure with --download-kokkos-kernels"); 295 PetscFunctionReturn(PETSC_SUCCESS); 296 } 297 298 PetscErrorCode FormJacobianLocalVec(DMDALocalInfo *info, Vec x, Mat jac, Mat jacpre, AppCtx *user) 299 { 300 PetscFunctionBeginUser; 301 PetscCheck(PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_SUP, "Need to reconfigure with --download-kokkos-kernels"); 302 PetscFunctionReturn(PETSC_SUCCESS); 303 } 304 #endif 305 306 /*TEST 307 308 build: 309 TODO: 310 311 TEST*/ 312