Actual source code: ex26.c
1: static const char help[] = "Test CUDA/HIP SF link reuse, cleanup, and completion on barrier contexts, and Kokkos VecScatter ordering with the current context.\n";
3: #include <petscsf.h>
4: #include <petscdevice.h>
5: #if PetscDefined(HAVE_CUDA)
6: #include <petscdevice_cuda.h>
7: #endif
8: #if PetscDefined(HAVE_HIP)
9: #include <petscdevice_hip.h>
10: #endif
12: static PetscErrorCode TestKokkosScatter(PetscDeviceContext saved, PetscDeviceContext dctx[])
13: {
14: Vec x, y;
15: VecScatter scatter;
16: IS from, to;
17: PetscDeviceContext current;
18: const PetscScalar *values;
19: PetscInt source[2], target[] = {0, 2};
20: PetscMPIInt rank, size;
22: PetscFunctionBeginUser;
23: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
24: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
25: // Initialize Kokkos on a context that stays alive until PetscFinalize().
26: PetscCall(PetscDeviceContextSetCurrentContext(saved));
27: PetscCall(VecCreate(PETSC_COMM_WORLD, &x));
28: PetscCall(VecSetSizes(x, 4, PETSC_DECIDE));
29: PetscCall(VecSetType(x, VECMPIKOKKOS));
30: PetscCall(VecCreate(PETSC_COMM_SELF, &y));
31: PetscCall(VecSetSizes(y, 4, 4));
32: PetscCall(VecSetType(y, VECSEQKOKKOS));
33: source[0] = 4 * rank + 3;
34: source[1] = 4 * ((rank + 1) % size) + 1;
35: PetscCall(ISCreateGeneral(PETSC_COMM_SELF, 2, source, PETSC_COPY_VALUES, &from));
36: PetscCall(ISCreateGeneral(PETSC_COMM_SELF, 2, target, PETSC_COPY_VALUES, &to));
37: PetscCall(VecScatterCreate(x, from, y, to, &scatter));
38: PetscCall(ISDestroy(&from));
39: PetscCall(ISDestroy(&to));
41: // Warm the cached links on Kokkos' stream, then reuse them on two different streams.
42: for (PetscInt pass = 0; pass < 3; ++pass) {
43: current = pass ? dctx[pass - 1] : saved;
44: PetscCall(PetscDeviceContextSetCurrentContext(current));
45: for (PetscInt local = 0; local < 2; ++local) {
46: for (PetscInt reverse = 0; reverse < 2; ++reverse) {
47: Vec input = reverse ? y : x, output = reverse ? x : y;
48: ScatterMode mode = reverse ? (local ? SCATTER_REVERSE_LOCAL : SCATTER_REVERSE) : (local ? SCATTER_FORWARD_LOCAL : SCATTER_FORWARD);
50: PetscCall(VecSet(input, 1));
51: PetscCall(VecSet(output, 0));
52: PetscCall(PetscDeviceContextSynchronize(saved));
53: PetscCall(PetscDeviceContextDelay(saved, 0.05));
54: PetscCall(VecScale(input, 2));
55: PetscCall(VecScatterBegin(scatter, input, output, INSERT_VALUES, mode));
56: PetscCall(PetscDeviceContextDelay(current, 0.05));
57: PetscCall(VecScatterEnd(scatter, input, output, INSERT_VALUES, mode));
58: PetscCall(VecScale(output, 3));
59: PetscCall(PetscDeviceContextSynchronize(saved));
60: PetscCall(PetscDeviceContextSynchronize(current));
61: PetscCall(VecGetArrayRead(output, &values));
62: for (PetscInt i = 0; i < 4; ++i) {
63: PetscScalar expected = 0;
65: if (i == (reverse ? 3 : 0) || (i == (reverse ? 1 : 2) && (!local || size == 1))) expected = 6;
66: PetscCheck(values[i] == expected, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Pass %" PetscInt_FMT ", mode %d: incorrect Kokkos scatter result at local entry %" PetscInt_FMT, pass, (int)mode, i);
67: }
68: PetscCall(VecRestoreArrayRead(output, &values));
69: }
70: }
71: }
72: PetscCall(VecScatterDestroy(&scatter));
73: PetscCall(VecDestroy(&x));
74: PetscCall(VecDestroy(&y));
75: PetscCall(PetscDeviceContextSetCurrentContext(saved));
76: PetscFunctionReturn(PETSC_SUCCESS);
77: }
79: int main(int argc, char **argv)
80: {
81: PetscSF sf;
82: PetscSFNode remote[2];
83: PetscDeviceContext saved, dctx[2], current;
84: PetscDevice device;
85: PetscDeviceType type;
86: PetscMemType mtype;
87: PetscStreamType streamtype = PETSC_STREAM_NONBLOCKING;
88: PetscInt local[] = {0, 2};
89: PetscInt *root, *leaf, *update, *input, *output;
90: PetscMPIInt rank, size;
91: PetscBool kokkos = PETSC_FALSE, cleanup = PETSC_FALSE;
93: PetscFunctionBeginUser;
94: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
95: PetscCall(PetscOptionsGetEnum(NULL, NULL, "-stream_type", PetscStreamTypes, (PetscEnum *)&streamtype, NULL));
96: PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_kokkos", &kokkos, NULL));
97: PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_cleanup", &cleanup, NULL));
98: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
99: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
100: PetscCall(PetscDeviceContextGetCurrentContext(&saved));
101: PetscCall(PetscDeviceContextGetDevice(saved, &device));
102: PetscCall(PetscDeviceGetType(device, &type));
103: mtype = type == PETSC_DEVICE_CUDA ? PETSC_MEMTYPE_CUDA : PETSC_MEMTYPE_HIP;
104: for (PetscInt i = 0; i < 2; ++i) {
105: PetscCall(PetscDeviceContextDuplicate(saved, &dctx[i]));
106: PetscCall(PetscDeviceContextSetStreamType(dctx[i], streamtype));
107: PetscCall(PetscDeviceContextSetUp(dctx[i]));
108: }
109: PetscCall(PetscDeviceContextSetCurrentContext(dctx[0]));
110: PetscCall(PetscDeviceMalloc(dctx[0], PETSC_MEMTYPE_DEVICE, 4, PETSC_DECIDE, &root));
111: PetscCall(PetscDeviceMalloc(dctx[0], PETSC_MEMTYPE_DEVICE, 4, PETSC_DECIDE, &leaf));
112: PetscCall(PetscDeviceMalloc(dctx[0], PETSC_MEMTYPE_DEVICE, 4, PETSC_DECIDE, &update));
113: PetscCall(PetscDeviceMalloc(dctx[0], PETSC_MEMTYPE_HOST, 4, PETSC_DECIDE, &input));
114: PetscCall(PetscDeviceMalloc(dctx[0], PETSC_MEMTYPE_HOST, 4, PETSC_DECIDE, &output));
115: PetscCall(PetscDeviceContextSynchronize(dctx[0]));
116: remote[0].rank = remote[1].rank = (rank + 1) % size;
117: remote[0].index = 3;
118: remote[1].index = 1;
119: PetscCall(PetscSFCreate(PETSC_COMM_WORLD, &sf));
120: PetscCall(PetscSFSetFromOptions(sf));
121: PetscCall(PetscSFSetGraph(sf, 4, 2, local, PETSC_COPY_VALUES, remote, PETSC_COPY_VALUES));
122: PetscCall(PetscSFSetUp(sf));
124: // Initialize the link on A, then reuse it on B and A. Complete each pass before switching contexts.
125: for (PetscInt pass = 0; pass < 3; ++pass) {
126: current = dctx[pass % 2];
127: PetscCall(PetscDeviceContextSetCurrentContext(current));
128: for (PetscInt i = 0; i < 4; ++i) input[i] = 100 * pass + 4 * rank + i;
129: for (PetscInt op = 0; op < 3; ++op) {
130: if (pass) PetscCall(PetscDeviceContextDelay(current, 0.05));
131: PetscCall(PetscDeviceMemcpy(current, root, input, 4 * sizeof(*root)));
132: PetscCall(PetscDeviceMemcpy(current, leaf, input, 4 * sizeof(*leaf)));
133: if (op == 0) {
134: PetscCall(PetscSFBcastWithMemTypeBegin(sf, MPIU_INT, mtype, root, mtype, leaf, MPI_REPLACE));
135: if (pass) PetscCall(PetscDeviceContextDelay(current, 0.05));
136: PetscCall(PetscSFBcastEnd(sf, MPIU_INT, root, leaf, MPI_REPLACE));
137: } else if (op == 1) {
138: PetscCall(PetscSFReduceWithMemTypeBegin(sf, MPIU_INT, mtype, leaf, mtype, root, MPI_SUM));
139: if (pass) PetscCall(PetscDeviceContextDelay(current, 0.05));
140: PetscCall(PetscSFReduceEnd(sf, MPIU_INT, leaf, root, MPI_SUM));
141: } else {
142: PetscCall(PetscSFFetchAndOpWithMemTypeBegin(sf, MPIU_INT, mtype, root, mtype, leaf, mtype, update, MPI_SUM));
143: if (pass) PetscCall(PetscDeviceContextDelay(current, 0.05));
144: PetscCall(PetscSFFetchAndOpEnd(sf, MPIU_INT, root, leaf, update, MPI_SUM));
145: }
146: if (streamtype == PETSC_STREAM_DEFAULT_WITH_BARRIER || streamtype == PETSC_STREAM_NONBLOCKING_WITH_BARRIER) {
147: PetscBool idle;
149: PetscCall(PetscDeviceContextQueryIdle(current, &idle));
150: PetscCheck(idle, PETSC_COMM_SELF, PETSC_ERR_PLIB, "SF operation %" PetscInt_FMT " returned with work pending on a barrier context", op);
151: }
152: if (op != 1) {
153: PetscCall(PetscDeviceMemcpy(current, output, op == 0 ? leaf : update, 4 * sizeof(*output)));
154: PetscCall(PetscDeviceContextSynchronize(current));
155: for (PetscInt i = 0; i < 2; ++i)
156: PetscCheck(output[local[i]] == 100 * pass + 4 * remote[i].rank + remote[i].index, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Pass %" PetscInt_FMT ", operation %" PetscInt_FMT ", incorrect leaf %" PetscInt_FMT, pass, op, local[i]);
157: }
158: if (op != 0) {
159: PetscMPIInt prev = (rank + size - 1) % size;
161: PetscCall(PetscDeviceMemcpy(current, output, root, 4 * sizeof(*output)));
162: PetscCall(PetscDeviceContextSynchronize(current));
163: for (PetscInt i = 0; i < 4; ++i) {
164: PetscInt expected = input[i];
166: if (i == 3) expected += 100 * pass + 4 * prev;
167: if (i == 1) expected += 100 * pass + 4 * prev + 2;
168: PetscCheck(output[i] == expected, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Pass %" PetscInt_FMT ", operation %" PetscInt_FMT ", incorrect root %" PetscInt_FMT, pass, op, i);
169: }
170: }
171: }
172: }
174: if (cleanup) {
175: for (PetscInt contiguous = 0; contiguous < 2; ++contiguous) {
176: PetscMPIInt prev = (rank + size - 1) % size;
178: // Contiguous indices isolate link buffer cleanup from the synchronous metadata cleanup.
179: if (contiguous) {
180: local[1] = 1;
181: remote[0].index = 1;
182: remote[1].index = 2;
183: PetscCall(PetscSFSetGraph(sf, 4, 2, local, PETSC_COPY_VALUES, remote, PETSC_COPY_VALUES));
184: }
185: PetscCall(PetscDeviceContextSetCurrentContext(dctx[0]));
186: for (PetscInt i = 0; i < 4; ++i) input[i] = 4 * rank + i + 1;
187: for (PetscInt pass = 0; pass < 2; ++pass) {
188: PetscCall(PetscDeviceMemcpy(dctx[0], root, input, 4 * sizeof(*root)));
189: PetscCall(PetscDeviceMemcpy(dctx[0], leaf, input, 4 * sizeof(*leaf)));
190: // Warm the link before delaying local scatter or remote unpacking.
191: if (pass && size == 1) PetscCall(PetscDeviceContextDelay(dctx[0], 0.05));
192: PetscCall(PetscSFReduceWithMemTypeBegin(sf, MPIU_INT, mtype, leaf, mtype, root, MPI_SUM));
193: if (pass && size > 1) PetscCall(PetscDeviceContextDelay(dctx[0], 0.05));
194: PetscCall(PetscSFReduceEnd(sf, MPIU_INT, leaf, root, MPI_SUM));
195: if (!pass) PetscCall(PetscDeviceContextSynchronize(dctx[0]));
196: }
197: // Cleanup must respect the link's stream even when a different context is current.
198: PetscCall(PetscDeviceContextSetCurrentContext(dctx[1]));
199: if (contiguous) PetscCall(PetscSFDestroy(&sf));
200: else PetscCall(PetscSFReset(sf));
201: PetscCall(PetscDeviceMemcpy(dctx[0], output, root, 4 * sizeof(*output)));
202: PetscCall(PetscDeviceContextSynchronize(dctx[0]));
203: for (PetscInt i = 0; i < 4; ++i) {
204: PetscInt expected = input[i];
206: for (PetscInt j = 0; j < 2; ++j)
207: if (i == remote[j].index) expected += 4 * prev + local[j] + 1;
208: PetscCheck(output[i] == expected, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Contiguous indices %" PetscInt_FMT ": incorrect root %" PetscInt_FMT " after SF cleanup", contiguous, i);
209: }
210: }
211: }
213: if (kokkos) PetscCall(TestKokkosScatter(saved, dctx));
214: PetscCall(PetscSFDestroy(&sf));
215: PetscCall(PetscDeviceFree(dctx[0], root));
216: PetscCall(PetscDeviceFree(dctx[0], leaf));
217: PetscCall(PetscDeviceFree(dctx[0], update));
218: PetscCall(PetscDeviceFree(dctx[0], input));
219: PetscCall(PetscDeviceFree(dctx[0], output));
220: PetscCall(PetscDeviceContextSynchronize(dctx[0]));
221: PetscCall(PetscDeviceContextSetCurrentContext(saved));
222: for (PetscInt i = 0; i < 2; ++i) PetscCall(PetscDeviceContextDestroy(&dctx[i]));
223: PetscCall(PetscFinalize());
224: return 0;
225: }
227: /*TEST
229: testset:
230: nsize: {{1 2}}
231: output_file: output/empty.out
232: args: -sf_type {{basic neighbor}} -use_gpu_aware_mpi 0 -stream_type {{nonblocking default_with_barrier nonblocking_with_barrier}} -test_cleanup
234: test:
235: suffix: cuda
236: requires: cuda
237: args: -sf_backend cuda
239: test:
240: suffix: hip
241: requires: hip
242: args: -sf_backend hip
244: test:
245: suffix: kokkos_cuda
246: requires: kokkos_kernels cuda
247: nsize: {{1 2}}
248: args: -test_kokkos -sf_type {{basic neighbor}} -sf_backend cuda -use_gpu_aware_mpi 0
249: output_file: output/empty.out
251: test:
252: suffix: kokkos_hip
253: requires: kokkos_kernels hip
254: nsize: {{1 2}}
255: args: -test_kokkos -sf_type {{basic neighbor}} -sf_backend hip -use_gpu_aware_mpi 0
256: output_file: output/empty.out
258: test:
259: suffix: cuda_nvshmem
260: requires: cuda nvshmem
261: nsize: 2
262: args: -sf_backend cuda -use_nvshmem -use_nvshmem_get {{0 1}}
263: output_file: output/empty.out
265: TEST*/