1 static char help[] = "Tests DMDAVecGetKokkosOffsetView() and DMDAVecGetKokkosOffsetViewDOF() \n\n";
2
3 #include <petscdmda_kokkos.hpp>
4 #include <petscdm.h>
5 #include <petscdmda.h>
6
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>;
12
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>;
15
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;
23
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>;
26
main(int argc,char ** argv)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 */
36
37 PetscFunctionBeginUser;
38 PetscCall(PetscInitialize(&argc, &argv, nullptr, help));
39
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;
47
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));
57
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));
66
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));
74
75 /* Scatter g, gg to l, ll */
76 PetscCall(DMGlobalToLocal(da, g, INSERT_VALUES, l));
77 PetscCall(DMGlobalToLocal(da, gg, INSERT_VALUES, ll));
78
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));
89
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));
98
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");
103
104 PetscCall(DMDestroy(&da));
105 PetscCall(VecDestroy(&l));
106 PetscCall(VecDestroy(&g));
107 PetscCall(VecDestroy(&ll));
108 PetscCall(VecDestroy(&gg));
109
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;
120
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));
130
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));
141
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));
152
153 /* Scatter g, gg to l, ll */
154 PetscCall(DMGlobalToLocal(da, g, INSERT_VALUES, l));
155 PetscCall(DMGlobalToLocal(da, gg, INSERT_VALUES, ll));
156
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));
169
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));
183
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");
188
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 }
197
198 /*TEST
199 build:
200 requires: kokkos_kernels
201
202 test:
203 suffix: 1
204 nsize: 4
205 requires: kokkos_kernels
206 args: -dm_vec_type kokkos
207 output_file: output/empty.out
208
209 TEST*/
210