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