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