Actual source code: ex302k.kokkos.cxx
1: static char help[] = "Testing MatCreateSeqAIJKokkosWithKokkosViews(), various read/write Kokkos view APIs on the device, and MatConvert() from MATMPIAIJ to MATAIJKOKKOS.\n\n";
3: #include <petsc_kokkos.hpp>
4: #include <petscvec_kokkos.hpp>
5: #include <petscdevice.h>
6: #include <petscmat.h>
7: #include <petscmat_kokkos.hpp>
8: #include <Kokkos_Core.hpp>
9: #include <Kokkos_DualView.hpp>
11: using HostMirrorMemorySpace = Kokkos::DualView<PetscScalar *>::host_mirror_space::memory_space;
13: static PetscErrorCode CheckBlockTypes(Mat A, MatType type)
14: {
15: Mat Ad, Ao;
16: PetscBool same;
18: PetscFunctionBeginUser;
19: PetscCall(MatMPIAIJGetSeqAIJ(A, &Ad, &Ao, NULL));
20: PetscCall(PetscObjectTypeCompare((PetscObject)Ad, type, &same));
21: PetscCheck(same, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Diagonal block is not %s", type);
22: PetscCall(PetscObjectTypeCompare((PetscObject)Ao, type, &same));
23: PetscCheck(same, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Off-diagonal block is not %s", type);
24: PetscFunctionReturn(PETSC_SUCCESS);
25: }
27: // MatConvert() from MATMPIAIJ to MATAIJKOKKOS must retype the blocks of the new matrix and leave those of the source alone
28: static PetscErrorCode TestConvertFromHost(Mat A)
29: {
30: Mat host, kok;
31: PetscBool equal;
33: PetscFunctionBeginUser;
34: PetscCall(MatConvert(A, MATMPIAIJ, MAT_INITIAL_MATRIX, &host));
35: PetscCall(CheckBlockTypes(host, MATSEQAIJ));
37: PetscCall(MatConvert(host, MATAIJKOKKOS, MAT_INITIAL_MATRIX, &kok));
38: PetscCall(CheckBlockTypes(kok, MATSEQAIJKOKKOS));
39: PetscCall(CheckBlockTypes(host, MATSEQAIJ));
40: PetscCall(MatMultEqual(host, kok, 3, &equal));
41: PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatConvert() with MAT_INITIAL_MATRIX gave wrong values");
43: PetscCall(MatScale(host, 2.0));
44: PetscCall(MatConvert(host, MATAIJKOKKOS, MAT_REUSE_MATRIX, &kok));
45: PetscCall(CheckBlockTypes(kok, MATSEQAIJKOKKOS));
46: PetscCall(CheckBlockTypes(host, MATSEQAIJ));
47: PetscCall(MatMultEqual(host, kok, 3, &equal));
48: PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatConvert() with MAT_REUSE_MATRIX gave wrong values");
50: PetscCall(MatDestroy(&kok));
51: PetscCall(MatDestroy(&host));
52: PetscFunctionReturn(PETSC_SUCCESS);
53: }
55: int main(int argc, char **argv)
56: {
57: Mat A, B, BT;
58: PetscInt i, j, column, M, N, m, n, m_ab, n_ab;
59: PetscInt *di, *dj, *oi, *oj, nd;
60: const PetscInt *garray;
61: PetscInt *garray_h;
62: PetscScalar *oa, *da;
63: PetscScalar value;
64: PetscRandom rctx;
65: PetscBool equal, done;
66: Mat AA, AB;
67: PetscMPIInt size, rank;
69: // ~~~~~~~~~~~~~~~~~~~~~
70: // This test shows the routines needed to build a kokkos matrix without preallocation
71: // on the host
72: // ~~~~~~~~~~~~~~~~~~~~~
74: PetscFunctionBeginUser;
75: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
76: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
77: PetscCheck(size > 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Must run with 2 or more processes");
78: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
80: /* Create a mpiaij matrix for checking */
81: PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, 5, 5, PETSC_DECIDE, PETSC_DECIDE, 0, NULL, 0, NULL, &A));
82: PetscCall(MatSetFromOptions(A));
83: PetscCall(MatSetOption(A, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_FALSE));
84: PetscCall(MatSetUp(A));
85: PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rctx));
86: PetscCall(PetscRandomSetFromOptions(rctx));
88: for (i = 5 * rank; i < 5 * rank + 5; i++) {
89: for (j = 0; j < 5 * size; j++) {
90: PetscCall(PetscRandomGetValue(rctx, &value));
91: column = (PetscInt)(5 * size * PetscRealPart(value));
93: // rank 0 has no off-process entries
94: if (rank == 0 && (column < i || column >= i)) column = i;
96: PetscCall(PetscRandomGetValue(rctx, &value));
97: PetscCall(MatSetValues(A, 1, &i, 1, &column, &value, INSERT_VALUES));
98: }
99: }
100: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
101: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
103: PetscCall(MatGetSize(A, &M, &N));
104: PetscCall(MatGetLocalSize(A, &m, &n));
105: PetscCall(MatMPIAIJGetSeqAIJ(A, &AA, &AB, &garray));
106: PetscCall(MatGetRowIJ(AA, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&di, (const PetscInt **)&dj, &done));
107: PetscCall(MatSeqAIJGetArray(AA, &da));
108: PetscCall(MatGetRowIJ(AB, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&oi, (const PetscInt **)&oj, &done));
109: PetscCall(MatSeqAIJGetArray(AB, &oa));
110: PetscCall(MatGetSize(AB, &m_ab, &n_ab));
112: Mat output_mat_local, output_mat_nonlocal;
113: // Be careful about scope given the kokkos memory reference counts
114: {
115: // Local
116: Kokkos::View<PetscScalar *> a_local_d;
117: Kokkos::View<PetscInt *> i_local_d;
118: Kokkos::View<PetscInt *> j_local_d;
120: // Nonlocal
121: Kokkos::View<PetscScalar *> a_nonlocal_d;
122: Kokkos::View<PetscInt *> i_nonlocal_d;
123: Kokkos::View<PetscInt *> j_nonlocal_d;
125: // Create device memory
126: PetscCallCXX(a_local_d = Kokkos::View<PetscScalar *>("a_local_d", di[5]));
127: PetscCallCXX(i_local_d = Kokkos::View<PetscInt *>("i_local_d", m + 1));
128: PetscCallCXX(j_local_d = Kokkos::View<PetscInt *>("j_local_d", di[5]));
130: // Create non-local device memory
131: PetscCallCXX(a_nonlocal_d = Kokkos::View<PetscScalar *>("a_nonlocal_d", oi[5]));
132: PetscCallCXX(i_nonlocal_d = Kokkos::View<PetscInt *>("i_nonlocal_d", m + 1));
133: PetscCallCXX(j_nonlocal_d = Kokkos::View<PetscInt *>("j_nonlocal_d", oi[5]));
135: // ~~~~~~~~~~~~~~~~~~~~~
136: // Could fill the aij on the device - we're just going to test
137: // by copying in the existing host values
138: // ~~~~~~~~~~~~~~~~~~~~~
139: Kokkos::View<PetscScalar *, HostMirrorMemorySpace> a_local_h;
140: Kokkos::View<PetscInt *, HostMirrorMemorySpace> i_local_h;
141: Kokkos::View<PetscInt *, HostMirrorMemorySpace> j_local_h;
142: Kokkos::View<PetscScalar *, HostMirrorMemorySpace> a_nonlocal_h;
143: Kokkos::View<PetscInt *, HostMirrorMemorySpace> i_nonlocal_h;
144: Kokkos::View<PetscInt *, HostMirrorMemorySpace> j_nonlocal_h;
146: PetscCallCXX(a_local_h = Kokkos::View<PetscScalar *, HostMirrorMemorySpace>(da, di[5]));
147: PetscCallCXX(i_local_h = Kokkos::View<PetscInt *, HostMirrorMemorySpace>(di, m + 1));
148: PetscCallCXX(j_local_h = Kokkos::View<PetscInt *, HostMirrorMemorySpace>(dj, di[5]));
149: PetscCallCXX(a_nonlocal_h = Kokkos::View<PetscScalar *, HostMirrorMemorySpace>(oa, oi[5]));
150: PetscCallCXX(i_nonlocal_h = Kokkos::View<PetscInt *, HostMirrorMemorySpace>(oi, m + 1));
151: PetscCallCXX(j_nonlocal_h = Kokkos::View<PetscInt *, HostMirrorMemorySpace>(oj, oi[5]));
153: // Haven't specified an exec space so these will all be synchronous
154: // and finish without a need to call fence after
155: PetscCallCXX(Kokkos::deep_copy(a_local_d, a_local_h));
156: PetscCallCXX(Kokkos::deep_copy(i_local_d, i_local_h));
157: PetscCallCXX(Kokkos::deep_copy(j_local_d, j_local_h));
158: PetscCallCXX(Kokkos::deep_copy(a_nonlocal_d, a_nonlocal_h));
159: PetscCallCXX(Kokkos::deep_copy(i_nonlocal_d, i_nonlocal_h));
160: PetscCallCXX(Kokkos::deep_copy(j_nonlocal_d, j_nonlocal_h));
162: // The garray passed in has to be on the host, but it can be created
163: // on device and copied to the host
164: // We're just going to copy the existing host values here
165: PetscCall(PetscMalloc1(n_ab, &garray_h));
166: for (int i = 0; i < n_ab; i++) garray_h[i] = garray[i];
168: // ~~~~~~~~~~~~~~~~~~~~~
170: // ~~~~~~~~~~~~~~~~~
171: // Test MatCreateSeqAIJKokkosWithKokkosViews
172: // ~~~~~~~~~~~~~~~~~
174: // We can create our local diagonal block matrix directly on the device
175: PetscCall(MatCreateSeqAIJKokkosWithKokkosViews(PETSC_COMM_SELF, m, n, i_local_d, j_local_d, a_local_d, &output_mat_local));
177: // We can create our nonlocal diagonal block matrix directly on the device
178: PetscCall(MatCreateSeqAIJKokkosWithKokkosViews(PETSC_COMM_SELF, m, n_ab, i_nonlocal_d, j_nonlocal_d, a_nonlocal_d, &output_mat_nonlocal));
180: // Build our MPI matrix
181: // If we provide garray and output_mat_nonlocal with local indices and the compactified size
182: // almost nothing happens on the host
183: PetscCall(MatCreateMPIAIJWithSeqAIJ(PETSC_COMM_WORLD, M, N, output_mat_local, output_mat_nonlocal, garray_h, &B));
185: PetscCall(MatEqual(A, B, &equal));
186: PetscCall(MatRestoreRowIJ(AA, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&di, (const PetscInt **)&dj, &done));
187: PetscCall(MatSeqAIJRestoreArray(AA, &da));
188: PetscCall(MatRestoreRowIJ(AB, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&oi, (const PetscInt **)&oj, &done));
189: PetscCall(MatSeqAIJRestoreArray(AB, &oa));
190: PetscCheck(equal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Likely a bug in MatCreateSeqAIJKokkosWithKokkosViews()");
192: // ~~~~~~~~~~~~~~~~~
193: // Test MatSeqAIJ{Get,Restore}KokkosView{,Write}
194: // ~~~~~~~~~~~~~~~~~
196: PetscReal norm_before, norm_after;
197: PetscInt nnz_AA = 0;
198: PetscScalar host_sum = 0.0;
200: PetscCall(MatNorm(AA, NORM_FROBENIUS, &norm_before));
202: /* Const overload: read-only access. Sum the values and check it matches a host-side reference. */
203: {
204: Kokkos::View<const PetscScalar *> ro_view;
205: PetscCall(MatSeqAIJGetKokkosView(AA, &ro_view));
206: nnz_AA = (PetscInt)ro_view.extent(0);
208: PetscScalar device_sum = 0.0;
209: Kokkos::parallel_reduce("ex302k_sum", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, nnz_AA), KOKKOS_LAMBDA(const PetscInt k, PetscScalar &lsum) { lsum += ro_view(k); }, device_sum);
211: const PetscScalar *aa_host;
212: PetscCall(MatSeqAIJGetArrayRead(AA, &aa_host));
213: for (PetscInt k = 0; k < nnz_AA; k++) host_sum += aa_host[k];
214: PetscCall(MatSeqAIJRestoreArrayRead(AA, &aa_host));
216: PetscCheck(PetscAbsScalar(device_sum - host_sum) <= PETSC_SMALL * PetscAbsScalar(host_sum), PETSC_COMM_SELF, PETSC_ERR_PLIB, "MatSeqAIJGetKokkosView (const) device/host mismatch");
217: PetscCall(MatSeqAIJRestoreKokkosView(AA, &ro_view));
218: }
220: /* Snapshot the original values into a host Kokkos view so we can restore them later via the write-only overload. */
221: Kokkos::View<PetscScalar *, HostMirrorMemorySpace> saved_h("saved_h", nnz_AA);
222: {
223: const PetscScalar *aa_host;
224: PetscCall(MatSeqAIJGetArrayRead(AA, &aa_host));
225: for (PetscInt k = 0; k < nnz_AA; k++) saved_h(k) = aa_host[k];
226: PetscCall(MatSeqAIJRestoreArrayRead(AA, &aa_host));
227: }
229: /* Write-only overload: overwrite all entries with zero (no read of the prior contents). */
230: {
231: Kokkos::View<PetscScalar *> rw_view;
232: PetscCall(MatSeqAIJGetKokkosViewWrite(AA, &rw_view));
233: PetscCallCXX(Kokkos::deep_copy(rw_view, (PetscScalar)0.0));
234: PetscCall(MatSeqAIJRestoreKokkosViewWrite(AA, &rw_view));
235: }
236: {
237: PetscReal norm_zero;
238: PetscCall(MatNorm(AA, NORM_FROBENIUS, &norm_zero));
239: PetscCheck(norm_zero == 0.0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "MatSeqAIJ{Get,Restore}KokkosViewWrite did not propagate device write");
240: }
242: /* Write-only overload again: restore the original values from the host snapshot. */
243: {
244: Kokkos::View<PetscScalar *> rw_view;
245: PetscCall(MatSeqAIJGetKokkosViewWrite(AA, &rw_view));
246: PetscCallCXX(Kokkos::deep_copy(rw_view, saved_h));
247: PetscCall(MatSeqAIJRestoreKokkosViewWrite(AA, &rw_view));
248: }
250: /* Read-write overload: scale by 2, then by 1/2 to undo. */
251: {
252: Kokkos::View<PetscScalar *> rw_view;
253: PetscCall(MatSeqAIJGetKokkosView(AA, &rw_view));
254: // It seems with SYCL, the user code has to use the same execution space as PETSc's. MatNorm() will sync the execution space before accessing data on host
255: Kokkos::parallel_for("ex302k_scale_up", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, nnz_AA), KOKKOS_LAMBDA(const PetscInt k) { rw_view(k) *= (PetscScalar)2.0; });
256: PetscCall(MatSeqAIJRestoreKokkosView(AA, &rw_view));
257: }
258: {
259: Kokkos::View<PetscScalar *> rw_view;
260: PetscCall(MatSeqAIJGetKokkosView(AA, &rw_view));
261: Kokkos::parallel_for("ex302k_scale_down", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, nnz_AA), KOKKOS_LAMBDA(const PetscInt k) { rw_view(k) *= (PetscScalar)0.5; });
262: PetscCall(MatSeqAIJRestoreKokkosView(AA, &rw_view));
263: }
265: PetscCall(MatNorm(AA, NORM_FROBENIUS, &norm_after));
266: PetscCheck(PetscAbsReal(norm_after - norm_before) <= PETSC_SMALL * norm_before, PETSC_COMM_SELF, PETSC_ERR_PLIB, "MatSeqAIJ{Get,Restore}KokkosView round trip changed AA");
267: }
269: PetscCall(MatTranspose(B, MAT_INITIAL_MATRIX, &BT));
271: PetscCall(TestConvertFromHost(A));
273: /* Free spaces */
274: PetscCall(PetscRandomDestroy(&rctx));
275: PetscCall(MatDestroy(&A));
276: PetscCall(MatDestroy(&B));
277: PetscCall(MatDestroy(&BT));
278: PetscCall(PetscFinalize());
279: return 0;
280: }
282: /*TEST
283: build:
284: requires: kokkos_kernels
286: test:
287: nsize: 2
288: args: -mat_type aijkokkos
289: requires: kokkos_kernels
290: output_file: output/empty.out
292: TEST*/