1: static const char help[] = "Test PetscSF with MPI large count (more than 2 billion elements in messages)\n\n";

  3: #include <petscsf.h>

  5: int main(int argc, char **argv)
  6: {
  7:   PetscSF            sf;
  8:   PetscInt           i, nroots, nleaves;
  9:   PetscInt           n       = (1ULL << 31) + 1024; /* a little over 2G elements */
 10:   PetscSFNode       *iremote = NULL;
 11:   PetscMPIInt        rank, size;
 12:   char              *rootdata = NULL, *leafdata = NULL;
 13:   Vec                x, y;
 14:   VecScatter         vscat;
 15:   PetscInt           rstart, rend;
 16:   IS                 ix;
 17:   const PetscScalar *xv;

 19:   PetscFunctionBeginUser;
 20:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 21:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
 22:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 23:   PetscCheck(size == 2, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "The test can only run with two MPI ranks");

 25:   /* Test PetscSF */
 26:   PetscCall(PetscSFCreate(PETSC_COMM_WORLD, &sf));
 27:   PetscCall(PetscSFSetFromOptions(sf));

 29:   if (rank == 0) {
 30:     nroots  = n;
 31:     nleaves = 0;
 32:   } else {
 33:     nroots  = 0;
 34:     nleaves = n;
 35:     PetscCall(PetscMalloc1(nleaves, &iremote));
 36:     for (i = 0; i < nleaves; i++) {
 37:       iremote[i].rank  = 0;
 38:       iremote[i].index = i;
 39:     }
 40:   }
 41:   PetscCall(PetscSFSetGraph(sf, nroots, nleaves, NULL, PETSC_COPY_VALUES, iremote, PETSC_COPY_VALUES));
 42:   PetscCall(PetscMalloc2(nroots, &rootdata, nleaves, &leafdata));
 43:   if (rank == 0) {
 44:     memset(rootdata, 11, nroots);
 45:     rootdata[nroots - 1] = 12; /* Use a different value at the end */
 46:   }

 48:   PetscCall(PetscSFBcastBegin(sf, MPI_SIGNED_CHAR, rootdata, leafdata, MPI_REPLACE)); /* rank 0->1, bcast rootdata to leafdata */
 49:   PetscCall(PetscSFBcastEnd(sf, MPI_SIGNED_CHAR, rootdata, leafdata, MPI_REPLACE));
 50:   PetscCall(PetscSFReduceBegin(sf, MPI_SIGNED_CHAR, leafdata, rootdata, MPI_SUM)); /* rank 1->0, add leafdata to rootdata */
 51:   PetscCall(PetscSFReduceEnd(sf, MPI_SIGNED_CHAR, leafdata, rootdata, MPI_SUM));
 52:   PetscCheck(rank != 0 || (rootdata[0] == 22 && rootdata[nroots - 1] == 24), PETSC_COMM_SELF, PETSC_ERR_PLIB, "SF: wrong results");

 54:   PetscCall(PetscFree2(rootdata, leafdata));
 55:   PetscCall(PetscFree(iremote));
 56:   PetscCall(PetscSFDestroy(&sf));

 58:   /* Test VecScatter */
 59:   PetscCall(VecCreate(PETSC_COMM_WORLD, &x));
 60:   PetscCall(VecCreate(PETSC_COMM_WORLD, &y));
 61:   PetscCall(VecSetSizes(x, rank == 0 ? n : 64, PETSC_DECIDE));
 62:   PetscCall(VecSetSizes(y, rank == 0 ? 64 : n, PETSC_DECIDE));
 63:   PetscCall(VecSetFromOptions(x));
 64:   PetscCall(VecSetFromOptions(y));

 66:   PetscCall(VecGetOwnershipRange(x, &rstart, &rend));
 67:   PetscCall(ISCreateStride(PETSC_COMM_SELF, rend - rstart, rstart, 1, &ix));
 68:   PetscCall(VecScatterCreate(x, ix, y, ix, &vscat));

 70:   PetscCall(VecSet(x, 3.0));
 71:   PetscCall(VecScatterBegin(vscat, x, y, INSERT_VALUES, SCATTER_FORWARD));
 72:   PetscCall(VecScatterEnd(vscat, x, y, INSERT_VALUES, SCATTER_FORWARD));

 74:   PetscCall(VecScatterBegin(vscat, y, x, ADD_VALUES, SCATTER_REVERSE));
 75:   PetscCall(VecScatterEnd(vscat, y, x, ADD_VALUES, SCATTER_REVERSE));

 77:   PetscCall(VecGetArrayRead(x, &xv));
 78:   PetscCheck(xv[0] == 6.0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "VecScatter: wrong results");
 79:   PetscCall(VecRestoreArrayRead(x, &xv));
 80:   PetscCall(VecDestroy(&x));
 81:   PetscCall(VecDestroy(&y));
 82:   PetscCall(VecScatterDestroy(&vscat));
 83:   PetscCall(ISDestroy(&ix));

 85:   PetscCall(PetscFinalize());
 86:   return 0;
 87: }

 89: /**TEST
 90:    test:
 91:      requires: defined(PETSC_HAVE_MPI_LARGE_COUNT) defined(PETSC_USE_64BIT_INDICES)
 92:      TODO: need a machine with big memory (~150GB) to run the test
 93:      nsize: 2
 94:      args: -sf_type {{basic neighbor}}

 96: TEST**/