1 static const char help[] = "Tests PetscDeviceContextFork/Join.\n\n"; 2 3 #include "petscdevicetestcommon.h" 4 5 static PetscErrorCode DoFork(PetscDeviceContext parent, PetscInt n, PetscDeviceContext **sub) { 6 PetscDeviceType dtype; 7 PetscStreamType stype; 8 9 PetscFunctionBegin; 10 PetscCall(AssertDeviceContextExists(parent)); 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(0); 26 } 27 28 static PetscErrorCode TestNestedPetscDeviceContextForkJoin(PetscDeviceContext parCtx, PetscDeviceContext *sub) { 29 const PetscInt nsub = 4; 30 PetscDeviceContext *subsub; 31 32 PetscFunctionBegin; 33 PetscValidDeviceContext(parCtx, 1); 34 PetscValidPointer(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(0); 44 } 45 46 /* test fork-join */ 47 static PetscErrorCode TestPetscDeviceContextForkJoin(PetscDeviceContext dctx) { 48 PetscDeviceContext *sub; 49 const PetscInt n = 10; 50 51 PetscFunctionBegin; 52 PetscValidDeviceContext(dctx, 1); 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(0); 69 } 70 71 int main(int argc, char *argv[]) { 72 MPI_Comm comm; 73 PetscDeviceContext dctx; 74 75 PetscFunctionBeginUser; 76 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 77 comm = PETSC_COMM_WORLD; 78 79 PetscCall(PetscDeviceContextCreate(&dctx)); 80 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)dctx, "local_")); 81 PetscCall(PetscDeviceContextSetFromOptions(comm, dctx)); 82 PetscCall(TestPetscDeviceContextForkJoin(dctx)); 83 PetscCall(PetscDeviceContextDestroy(&dctx)); 84 85 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 86 PetscCall(TestPetscDeviceContextForkJoin(dctx)); 87 88 PetscCall(PetscPrintf(comm, "EXIT_SUCCESS\n")); 89 PetscCall(PetscFinalize()); 90 return 0; 91 } 92 93 /*TEST 94 95 build: 96 requires: defined(PETSC_HAVE_CXX) 97 98 testset: 99 output_file: ./output/ExitSuccess.out 100 nsize: {{1 3}} 101 args: -device_enable {{lazy eager}} 102 args: -local_device_context_stream_type {{global_blocking default_blocking global_nonblocking}} 103 test: 104 requires: !device 105 suffix: host_no_device 106 test: 107 requires: device 108 args: -root_device_context_device_type host 109 suffix: host_with_device 110 test: 111 requires: cuda 112 args: -root_device_context_device_type cuda 113 suffix: cuda 114 test: 115 requires: hip 116 args: -root_device_context_device_type hip 117 suffix: hip 118 119 TEST*/ 120