Actual source code: ex1k.kokkos.cxx

  1: static char help[] = "Tests DMDAVecGetKokkosOffsetView() and DMDAVecGetKokkosOffsetViewDOF() \n\n";

  3: #include <petscdmda_kokkos.hpp>
  4: #include <petscdm.h>
  5: #include <petscdmda.h>

  7: using Kokkos::Iterate;
  8: using Kokkos::MDRangePolicy;
  9: using Kokkos::Rank;
 10: using PetscScalarKokkosOffsetView2D      = Kokkos::Experimental::OffsetView<PetscScalar **, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::memory_space>;
 11: using ConstPetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<const PetscScalar **, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::memory_space>;

 13: using PetscScalarKokkosOffsetView3D      = Kokkos::Experimental::OffsetView<PetscScalar ***, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::memory_space>;
 14: using ConstPetscScalarKokkosOffsetView3D = Kokkos::Experimental::OffsetView<const PetscScalar ***, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::memory_space>;

 16: /* can not define the type inside main, otherwise have this compilation error:
 17:    error: A type local to a function ("Node") cannot be used in the type of a
 18:    variable captured by an extended __device__ or __host__ __device__ lambda
 19: */
 20: typedef struct {
 21:   PetscScalar u, v, w;
 22: } Node;

 24: using NodeKokkosOffsetView2D      = Kokkos::Experimental::OffsetView<Node **, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::memory_space>;
 25: using ConstNodeKokkosOffsetView2D = Kokkos::Experimental::OffsetView<const Node **, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::memory_space>;

 27: int main(int argc, char **argv)
 28: {
 29:   DM              da;
 30:   PetscInt        M = 5, N = 7, xm, ym, xs, ys;
 31:   PetscInt        dof = 1, sw = 1;
 32:   DMBoundaryType  bx = DM_BOUNDARY_PERIODIC, by = DM_BOUNDARY_PERIODIC;
 33:   DMDAStencilType st = DMDA_STENCIL_STAR;
 34:   PetscReal       nrm;
 35:   Vec             g, l, gg, ll; /* global/local vectors of the da */

 38:   PetscInitialize(&argc, &argv, (char *)0, help);

 40:   /* ===========================================================================
 41:     Show how to manage a multi-component DMDA with DMDAVecGetKokkosOffsetViewDOF
 42:    ============================================================================*/
 43:   PetscScalar                     ***garray; /* arrays of g, ll */
 44:   const PetscScalar               ***larray;
 45:   PetscScalarKokkosOffsetView3D      gview; /* views of gg, ll */
 46:   ConstPetscScalarKokkosOffsetView3D lview;

 48:   dof = 2;
 49:   DMDACreate2d(PETSC_COMM_WORLD, bx, by, st, M, N, PETSC_DECIDE, PETSC_DECIDE, dof, sw, NULL, NULL, &da);
 50:   DMSetFromOptions(da);
 51:   DMSetUp(da);
 52:   DMDAGetCorners(da, &xs, &ys, 0, &xm, &ym, 0);
 53:   DMCreateGlobalVector(da, &g);
 54:   DMCreateLocalVector(da, &l);
 55:   DMCreateGlobalVector(da, &gg);
 56:   DMCreateLocalVector(da, &ll);

 58:   /* Init g using array */
 59:   DMDAVecGetArrayDOFWrite(da, g, &garray);
 60:   for (PetscInt j = ys; j < ys + ym; j++) { /* run on host */
 61:     for (PetscInt i = xs; i < xs + xm; i++) {
 62:       for (PetscInt c = 0; c < dof; c++) garray[j][i][c] = 100 * j + 10 * (i + 1) + c;
 63:     }
 64:   }
 65:   DMDAVecRestoreArrayDOFWrite(da, g, &garray);

 67:   /* Init gg using view */
 68:   DMDAVecGetKokkosOffsetViewDOFWrite(da, gg, &gview);
 69:   Kokkos::parallel_for(
 70:     "init 1", MDRangePolicy<Rank<3, Iterate::Right, Iterate::Right>>({ys, xs, 0}, {ys + ym, xs + xm, dof}),
 71:     KOKKOS_LAMBDA(PetscInt j, PetscInt i, PetscInt c) /* might run on device */
 72:     { gview(j, i, c) = 100 * j + 10 * (i + 1) + c; });
 73:   DMDAVecRestoreKokkosOffsetViewDOFWrite(da, gg, &gview);

 75:   /* Scatter g, gg to l, ll */
 76:   DMGlobalToLocal(da, g, INSERT_VALUES, l);
 77:   DMGlobalToLocal(da, gg, INSERT_VALUES, ll);

 79:   /* Do stencil on g with values from l that contains ghosts */
 80:   DMDAVecGetArrayDOFWrite(da, g, &garray);
 81:   DMDAVecGetArrayDOFRead(da, l, &larray);
 82:   for (PetscInt j = ys; j < ys + ym; j++) {
 83:     for (PetscInt i = xs; i < xs + xm; i++) {
 84:       for (PetscInt c = 0; c < dof; c++) garray[j][i][c] = (larray[j][i - 1][c] + larray[j][i + 1][c] + larray[j - 1][i][c] + larray[j + 1][i][c]) / 4.0;
 85:     }
 86:   }
 87:   DMDAVecRestoreArrayDOFWrite(da, g, &garray);
 88:   DMDAVecRestoreArrayDOFRead(da, l, &larray);

 90:   /* Do stencil on gg with values from ll that contains ghosts */
 91:   DMDAVecGetKokkosOffsetViewDOFWrite(da, gg, &gview);
 92:   DMDAVecGetKokkosOffsetViewDOF(da, ll, &lview);
 93:   Kokkos::parallel_for(
 94:     "stencil 1", MDRangePolicy<Rank<3, Iterate::Right, Iterate::Right>>({ys, xs, 0}, {ys + ym, xs + xm, dof}),
 95:     KOKKOS_LAMBDA(PetscInt j, PetscInt i, PetscInt c) { gview(j, i, c) = (lview(j, i - 1, c) + lview(j, i + 1, c) + lview(j - 1, i, c) + lview(j + 1, i, c)) / 4.0; });
 96:   DMDAVecRestoreKokkosOffsetViewDOFWrite(da, gg, &gview);
 97:   DMDAVecRestoreKokkosOffsetViewDOF(da, ll, &lview);

 99:   /* gg should be equal to g */
100:   VecAXPY(g, -1.0, gg);
101:   VecNorm(g, NORM_2, &nrm);

104:   DMDestroy(&da);
105:   VecDestroy(&l);
106:   VecDestroy(&g);
107:   VecDestroy(&ll);
108:   VecDestroy(&gg);

110:   /* =============================================================================
111:     Show how to manage a multi-component DMDA using DMDAVecGetKokkosOffsetView and
112:     a customized struct type
113:    ==============================================================================*/
114:   Node                             **garray2; /* Node arrays of g, l */
115:   const Node                       **larray2;
116:   PetscScalarKokkosOffsetView2D      gview2; /* PetscScalar views of gg, ll */
117:   ConstPetscScalarKokkosOffsetView2D lview2;
118:   NodeKokkosOffsetView2D             gnview; /* Node views of gg, ll */
119:   ConstNodeKokkosOffsetView2D        lnview;

121:   dof = sizeof(Node) / sizeof(PetscScalar);
122:   DMDACreate2d(PETSC_COMM_WORLD, bx, by, st, M, N, PETSC_DECIDE, PETSC_DECIDE, sizeof(Node) / sizeof(PetscScalar), sw, NULL, NULL, &da);
123:   DMSetFromOptions(da);
124:   DMSetUp(da);
125:   DMDAGetCorners(da, &xs, &ys, 0, &xm, &ym, 0);
126:   DMCreateGlobalVector(da, &g);
127:   DMCreateLocalVector(da, &l);
128:   DMCreateGlobalVector(da, &gg);
129:   DMCreateLocalVector(da, &ll);

131:   /* Init g using array */
132:   DMDAVecGetArrayWrite(da, g, &garray2);
133:   for (PetscInt j = ys; j < ys + ym; j++) {
134:     for (PetscInt i = xs; i < xs + xm; i++) {
135:       garray2[j][i].u = 100 * j + 10 * (i + 1) + 111;
136:       garray2[j][i].v = 100 * j + 10 * (i + 1) + 222;
137:       garray2[j][i].w = 100 * j + 10 * (i + 1) + 333;
138:     }
139:   }
140:   DMDAVecRestoreArrayWrite(da, g, &garray2);

142:   /* Init gg using view */
143:   DMDAVecGetKokkosOffsetViewWrite(da, gg, &gview2);
144:   gnview = NodeKokkosOffsetView2D(reinterpret_cast<Node *>(gview2.data()), {gview2.begin(0) / dof, gview2.begin(1) / dof}, {gview2.end(0) / dof, gview2.end(1) / dof});
145:   Kokkos::parallel_for(
146:     "init 2", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscInt j, PetscInt i) {
147:       gnview(j, i).u = 100 * j + 10 * (i + 1) + 111;
148:       gnview(j, i).v = 100 * j + 10 * (i + 1) + 222;
149:       gnview(j, i).w = 100 * j + 10 * (i + 1) + 333;
150:     });
151:   DMDAVecRestoreKokkosOffsetViewWrite(da, gg, &gview2);

153:   /* Scatter g, gg to l, ll */
154:   DMGlobalToLocal(da, g, INSERT_VALUES, l);
155:   DMGlobalToLocal(da, gg, INSERT_VALUES, ll);

157:   /* Do stencil on g with values from l that contains ghosts */
158:   DMDAVecGetArrayWrite(da, g, &garray2);
159:   DMDAVecGetArray(da, l, &larray2);
160:   for (PetscInt j = ys; j < ys + ym; j++) {
161:     for (PetscInt i = xs; i < xs + xm; i++) {
162:       garray2[j][i].u = (larray2[j][i - 1].u + larray2[j][i + 1].u + larray2[j - 1][i].u + larray2[j + 1][i].u) / 4.0;
163:       garray2[j][i].v = (larray2[j][i - 1].v + larray2[j][i + 1].v + larray2[j - 1][i].v + larray2[j + 1][i].v) / 4.0;
164:       garray2[j][i].w = (larray2[j][i - 1].w + larray2[j][i + 1].w + larray2[j - 1][i].w + larray2[j + 1][i].w) / 4.0;
165:     }
166:   }
167:   DMDAVecRestoreArrayWrite(da, g, &garray2);
168:   DMDAVecRestoreArray(da, l, &larray2);

170:   /* Do stencil on gg with values from ll that contains ghosts */
171:   DMDAVecGetKokkosOffsetViewWrite(da, gg, &gview2); /* write-only */
172:   DMDAVecGetKokkosOffsetView(da, ll, &lview2);      /* read-only */
173:   gnview = NodeKokkosOffsetView2D(reinterpret_cast<Node *>(gview2.data()), {gview2.begin(0) / dof, gview2.begin(1) / dof}, {gview2.end(0) / dof, gview2.end(1) / dof});
174:   lnview = ConstNodeKokkosOffsetView2D(reinterpret_cast<const Node *>(lview2.data()), {lview2.begin(0) / dof, lview2.begin(1) / dof}, {lview2.end(0) / dof, lview2.end(1) / dof});
175:   Kokkos::parallel_for(
176:     "stencil 2", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscInt j, PetscInt i) {
177:       gnview(j, i).u = (lnview(j, i - 1).u + lnview(j, i + 1).u + lnview(j - 1, i).u + lnview(j + 1, i).u) / 4.0;
178:       gnview(j, i).v = (lnview(j, i - 1).v + lnview(j, i + 1).v + lnview(j - 1, i).v + lnview(j + 1, i).v) / 4.0;
179:       gnview(j, i).w = (lnview(j, i - 1).w + lnview(j, i + 1).w + lnview(j - 1, i).w + lnview(j + 1, i).w) / 4.0;
180:     });
181:   DMDAVecRestoreKokkosOffsetViewWrite(da, gg, &gview2);
182:   DMDAVecRestoreKokkosOffsetView(da, ll, &lview2);

184:   /* gg should be equal to g */
185:   VecAXPY(g, -1.0, gg);
186:   VecNorm(g, NORM_2, &nrm);

189:   DMDestroy(&da);
190:   VecDestroy(&l);
191:   VecDestroy(&g);
192:   VecDestroy(&ll);
193:   VecDestroy(&gg);
194:   PetscFinalize();
195:   return 0;
196: }

198: /*TEST
199:   build:
200:     requires: kokkos_kernels

202:   test:
203:     suffix: 1
204:     nsize: 4
205:     requires: kokkos_kernels
206:     args: -dm_vec_type kokkos

208: TEST*/