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