1 static const char help[] = "Tests PetscDeviceContextFork/Join.\n\n";
2
3 #include "petscdevicetestcommon.h"
4
DoFork(PetscDeviceContext parent,PetscInt n,PetscDeviceContext ** sub)5 static PetscErrorCode DoFork(PetscDeviceContext parent, PetscInt n, PetscDeviceContext **sub)
6 {
7 PetscDeviceType dtype;
8 PetscStreamType stype;
9
10 PetscFunctionBegin;
11 PetscCall(PetscDeviceContextGetDeviceType(parent, &dtype));
12 PetscCall(PetscDeviceContextGetStreamType(parent, &stype));
13 PetscCall(PetscDeviceContextFork(parent, n, sub));
14 if (n) PetscCheck(*sub, PETSC_COMM_SELF, PETSC_ERR_PLIB, "PetscDeviceContextFork() return NULL pointer for %" PetscInt_FMT " children", n);
15 for (PetscInt i = 0; i < n; ++i) {
16 PetscDeviceType sub_dtype;
17 PetscStreamType sub_stype;
18
19 PetscCall(AssertDeviceContextExists((*sub)[i]));
20 PetscCall(PetscDeviceContextGetStreamType((*sub)[i], &sub_stype));
21 PetscCall(AssertPetscStreamTypesValidAndEqual(sub_stype, stype, "Child stream type %s != parent stream type %s"));
22 PetscCall(PetscDeviceContextGetDeviceType((*sub)[i], &sub_dtype));
23 PetscCall(AssertPetscDeviceTypesValidAndEqual(sub_dtype, dtype, "Child device type %s != parent device type %s"));
24 }
25 PetscFunctionReturn(PETSC_SUCCESS);
26 }
27
TestNestedPetscDeviceContextForkJoin(PetscDeviceContext parCtx,PetscDeviceContext * sub)28 static PetscErrorCode TestNestedPetscDeviceContextForkJoin(PetscDeviceContext parCtx, PetscDeviceContext *sub)
29 {
30 const PetscInt nsub = 4;
31 PetscDeviceContext *subsub;
32
33 PetscFunctionBegin;
34 PetscAssertPointer(sub, 2);
35 PetscCall(AssertPetscDeviceContextsValidAndEqual(parCtx, sub[0], "Current global context does not match expected global context"));
36 /* create some children from an active child */
37 PetscCall(DoFork(sub[1], nsub, &subsub));
38 /* join on a sibling to the parent */
39 PetscCall(PetscDeviceContextJoin(sub[2], nsub - 2, PETSC_DEVICE_CONTEXT_JOIN_SYNC, &subsub));
40 /* join on the grandparent */
41 PetscCall(PetscDeviceContextJoin(parCtx, nsub - 2, PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC, &subsub));
42 PetscCall(PetscDeviceContextJoin(sub[1], nsub, PETSC_DEVICE_CONTEXT_JOIN_DESTROY, &subsub));
43 PetscFunctionReturn(PETSC_SUCCESS);
44 }
45
46 /* test fork-join */
TestPetscDeviceContextForkJoin(PetscDeviceContext dctx)47 static PetscErrorCode TestPetscDeviceContextForkJoin(PetscDeviceContext dctx)
48 {
49 PetscDeviceContext *sub;
50 const PetscInt n = 10;
51
52 PetscFunctionBegin;
53 /* mostly for valgrind to catch errors */
54 PetscCall(DoFork(dctx, n, &sub));
55 PetscCall(PetscDeviceContextJoin(dctx, n, PETSC_DEVICE_CONTEXT_JOIN_DESTROY, &sub));
56 /* do it twice */
57 PetscCall(DoFork(dctx, n, &sub));
58 PetscCall(PetscDeviceContextJoin(dctx, n, PETSC_DEVICE_CONTEXT_JOIN_DESTROY, &sub));
59
60 /* create some children */
61 PetscCall(DoFork(dctx, n + 1, &sub));
62 /* test forking within nested function */
63 PetscCall(TestNestedPetscDeviceContextForkJoin(sub[0], sub));
64 /* join a subset */
65 PetscCall(PetscDeviceContextJoin(dctx, n - 1, PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC, &sub));
66 /* back to the ether from whence they came */
67 PetscCall(PetscDeviceContextJoin(dctx, n + 1, PETSC_DEVICE_CONTEXT_JOIN_DESTROY, &sub));
68 PetscFunctionReturn(PETSC_SUCCESS);
69 }
70
main(int argc,char * argv[])71 int main(int argc, char *argv[])
72 {
73 MPI_Comm comm;
74 PetscDeviceContext dctx;
75
76 PetscFunctionBeginUser;
77 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
78 comm = PETSC_COMM_WORLD;
79
80 PetscCall(PetscDeviceContextCreate(&dctx));
81 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)dctx, "local_"));
82 PetscCall(PetscDeviceContextSetFromOptions(comm, dctx));
83 PetscCall(TestPetscDeviceContextForkJoin(dctx));
84 PetscCall(PetscDeviceContextDestroy(&dctx));
85
86 PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
87 PetscCall(TestPetscDeviceContextForkJoin(dctx));
88
89 PetscCall(TestPetscDeviceContextForkJoin(NULL));
90
91 PetscCall(PetscPrintf(comm, "EXIT_SUCCESS\n"));
92 PetscCall(PetscFinalize());
93 return 0;
94 }
95
96 /*TEST
97
98 testset:
99 requires: defined(PETSC_DEVICELANGUAGE_CXX)
100 output_file: output/ExitSuccess.out
101 nsize: {{1 3}}
102 args: -device_enable {{lazy eager}}
103 args: -local_device_context_stream_type {{default nonblocking default_with_barrier nonblocking_with_barrier}}
104 test:
105 requires: !device
106 suffix: host_no_device
107 test:
108 requires: device
109 args: -root_device_context_device_type host
110 suffix: host_with_device
111 test:
112 requires: cuda
113 args: -root_device_context_device_type cuda
114 suffix: cuda
115 test:
116 requires: hip
117 args: -root_device_context_device_type hip
118 suffix: hip
119
120 TEST*/
121