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