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