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