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 */

 37:   PetscFunctionBeginUser;
 38:   PetscCall(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:   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, bx, by, st, M, N, PETSC_DECIDE, PETSC_DECIDE, dof, sw, NULL, NULL, &da));
 50:   PetscCall(DMSetFromOptions(da));
 51:   PetscCall(DMSetUp(da));
 52:   PetscCall(DMDAGetCorners(da, &xs, &ys, 0, &xm, &ym, 0));
 53:   PetscCall(DMCreateGlobalVector(da, &g));
 54:   PetscCall(DMCreateLocalVector(da, &l));
 55:   PetscCall(DMCreateGlobalVector(da, &gg));
 56:   PetscCall(DMCreateLocalVector(da, &ll));

 58:   /* Init g using array */
 59:   PetscCall(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:   PetscCall(DMDAVecRestoreArrayDOFWrite(da, g, &garray));

 67:   /* Init gg using view */
 68:   PetscCall(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:   PetscCall(DMDAVecRestoreKokkosOffsetViewDOFWrite(da, gg, &gview));

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

 79:   /* Do stencil on g with values from l that contains ghosts */
 80:   PetscCall(DMDAVecGetArrayDOFWrite(da, g, &garray));
 81:   PetscCall(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:   PetscCall(DMDAVecRestoreArrayDOFWrite(da, g, &garray));
 88:   PetscCall(DMDAVecRestoreArrayDOFRead(da, l, &larray));

 90:   /* Do stencil on gg with values from ll that contains ghosts */
 91:   PetscCall(DMDAVecGetKokkosOffsetViewDOFWrite(da, gg, &gview));
 92:   PetscCall(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:   PetscCall(DMDAVecRestoreKokkosOffsetViewDOFWrite(da, gg, &gview));
 97:   PetscCall(DMDAVecRestoreKokkosOffsetViewDOF(da, ll, &lview));

 99:   /* gg should be equal to g */
100:   PetscCall(VecAXPY(g, -1.0, gg));
101:   PetscCall(VecNorm(g, NORM_2, &nrm));
102:   PetscCheck(nrm < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "gg is not equal to g");

104:   PetscCall(DMDestroy(&da));
105:   PetscCall(VecDestroy(&l));
106:   PetscCall(VecDestroy(&g));
107:   PetscCall(VecDestroy(&ll));
108:   PetscCall(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:   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, bx, by, st, M, N, PETSC_DECIDE, PETSC_DECIDE, sizeof(Node) / sizeof(PetscScalar), sw, NULL, NULL, &da));
123:   PetscCall(DMSetFromOptions(da));
124:   PetscCall(DMSetUp(da));
125:   PetscCall(DMDAGetCorners(da, &xs, &ys, 0, &xm, &ym, 0));
126:   PetscCall(DMCreateGlobalVector(da, &g));
127:   PetscCall(DMCreateLocalVector(da, &l));
128:   PetscCall(DMCreateGlobalVector(da, &gg));
129:   PetscCall(DMCreateLocalVector(da, &ll));

131:   /* Init g using array */
132:   PetscCall(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:   PetscCall(DMDAVecRestoreArrayWrite(da, g, &garray2));

142:   /* Init gg using view */
143:   PetscCall(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:   PetscCall(DMDAVecRestoreKokkosOffsetViewWrite(da, gg, &gview2));

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

157:   /* Do stencil on g with values from l that contains ghosts */
158:   PetscCall(DMDAVecGetArrayWrite(da, g, &garray2));
159:   PetscCall(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:   PetscCall(DMDAVecRestoreArrayWrite(da, g, &garray2));
168:   PetscCall(DMDAVecRestoreArray(da, l, &larray2));

170:   /* Do stencil on gg with values from ll that contains ghosts */
171:   PetscCall(DMDAVecGetKokkosOffsetViewWrite(da, gg, &gview2)); /* write-only */
172:   PetscCall(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:   PetscCall(DMDAVecRestoreKokkosOffsetViewWrite(da, gg, &gview2));
182:   PetscCall(DMDAVecRestoreKokkosOffsetView(da, ll, &lview2));

184:   /* gg should be equal to g */
185:   PetscCall(VecAXPY(g, -1.0, gg));
186:   PetscCall(VecNorm(g, NORM_2, &nrm));
187:   PetscCheck(nrm < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "gg is not equal to g");

189:   PetscCall(DMDestroy(&da));
190:   PetscCall(VecDestroy(&l));
191:   PetscCall(VecDestroy(&g));
192:   PetscCall(VecDestroy(&ll));
193:   PetscCall(VecDestroy(&gg));
194:   PetscCall(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*/