Actual source code: ex4.c

  1: static const char help[] = "Tests PetscDeviceContextFork/Join.\n\n";

  3: #include "petscdevicetestcommon.h"

  5: static PetscErrorCode DoFork(PetscDeviceContext parent, PetscInt n, PetscDeviceContext **sub)
  6: {
  7:   PetscDeviceType dtype;
  8:   PetscStreamType stype;

 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;

 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: }

 28: static PetscErrorCode TestNestedPetscDeviceContextForkJoin(PetscDeviceContext parCtx, PetscDeviceContext *sub)
 29: {
 30:   const PetscInt      nsub = 4;
 31:   PetscDeviceContext *subsub;

 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: }

 46: /* test fork-join */
 47: static PetscErrorCode TestPetscDeviceContextForkJoin(PetscDeviceContext dctx)
 48: {
 49:   PetscDeviceContext *sub;
 50:   const PetscInt      n = 10;

 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));

 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: }

 71: int main(int argc, char *argv[])
 72: {
 73:   MPI_Comm           comm;
 74:   PetscDeviceContext dctx;

 76:   PetscFunctionBeginUser;
 77:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 78:   comm = PETSC_COMM_WORLD;

 80:   PetscCall(PetscDeviceContextCreate(&dctx));
 81:   PetscCall(PetscObjectSetOptionsPrefix((PetscObject)dctx, "local_"));
 82:   PetscCall(PetscDeviceContextSetFromOptions(comm, dctx));
 83:   PetscCall(TestPetscDeviceContextForkJoin(dctx));
 84:   PetscCall(PetscDeviceContextDestroy(&dctx));

 86:   PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
 87:   PetscCall(TestPetscDeviceContextForkJoin(dctx));

 89:   PetscCall(TestPetscDeviceContextForkJoin(NULL));

 91:   PetscCall(PetscPrintf(comm, "EXIT_SUCCESS\n"));
 92:   PetscCall(PetscFinalize());
 93:   return 0;
 94: }

 96: /*TEST

 98:   testset:
 99:     requires: 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

120: TEST*/