Actual source code: ex22.c

  1: static const char help[] = "Test PetscSFFetchAndOp \n\n";

  3: #include <petscvec.h>
  4: #include <petscsf.h>

  6: int main(int argc, char *argv[])
  7: {
  8:   PetscInt           n, N = 12;
  9:   PetscInt          *indices;
 10:   IS                 ix, iy;
 11:   VecScatter         vscat;
 12:   Vec                x, y, z;
 13:   PetscInt           rstart, rend;
 14:   const PetscScalar *xarray;
 15:   PetscScalar       *yarray, *zarray;
 16:   PetscMemType       xmtype, ymtype, zmtype;

 18:   PetscFunctionBeginUser;
 19:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));

 21:   PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, PETSC_DECIDE, N, &x));
 22:   PetscCall(VecDuplicate(x, &y));
 23:   PetscCall(VecDuplicate(x, &z));
 24:   PetscCall(VecGetLocalSize(x, &n));

 26:   PetscCall(VecGetOwnershipRange(x, &rstart, &rend));
 27:   PetscCall(ISCreateStride(PETSC_COMM_WORLD, n, rstart, 1, &ix));
 28:   PetscCall(PetscMalloc1(n, &indices));
 29:   for (PetscInt i = rstart; i < rend; i++) indices[i - rstart] = i / 4;
 30:   PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, n, indices, PETSC_OWN_POINTER, &iy));

 32:   // connect y[0] to x[0..3], y[1] to x[4..7], etc
 33:   PetscCall(VecScatterCreate(y, iy, x, ix, &vscat)); // y has roots, x has leaves

 35:   PetscCall(VecSet(x, 1.0));
 36:   PetscCall(VecSet(y, 2.0));

 38:   PetscCall(VecGetArrayReadAndMemType(x, &xarray, &xmtype));
 39:   PetscCall(VecGetArrayAndMemType(y, &yarray, &ymtype));
 40:   PetscCall(VecGetArrayWriteAndMemType(z, &zarray, &zmtype));

 42:   PetscCall(PetscSFFetchAndOpWithMemTypeBegin(vscat, MPIU_SCALAR, ymtype, yarray, xmtype, xarray, zmtype, zarray, MPI_SUM));
 43:   PetscCall(PetscSFFetchAndOpEnd(vscat, MPIU_SCALAR, yarray, xarray, zarray, MPI_SUM));

 45:   PetscCall(VecRestoreArrayReadAndMemType(x, &xarray));
 46:   PetscCall(VecRestoreArrayAndMemType(y, &yarray));
 47:   PetscCall(VecRestoreArrayWriteAndMemType(z, &zarray));

 49:   PetscCall(VecView(y, PETSC_VIEWER_STDOUT_WORLD));
 50:   PetscCall(VecView(z, PETSC_VIEWER_STDOUT_WORLD));
 51:   PetscCall(ISDestroy(&ix));
 52:   PetscCall(ISDestroy(&iy));
 53:   PetscCall(VecDestroy(&x));
 54:   PetscCall(VecDestroy(&y));
 55:   PetscCall(VecDestroy(&z));
 56:   PetscCall(VecScatterDestroy(&vscat));
 57:   PetscCall(PetscFinalize());
 58: }

 60: /*TEST
 61:   testset:
 62:     nsize: {{1 4}}
 63:     # since FetchAndOp on complex would need to be atomic in this test
 64:     requires: !complex
 65:     output_file: output/ex22.out
 66:     filter: grep -v "type" | grep -v "Process" |grep -v "Vec Object"

 68:     test:
 69:       suffix: cuda
 70:       requires: cuda
 71:       args: -vec_type cuda

 73:     test:
 74:       suffix: hip
 75:       requires: hip
 76:       args: -vec_type hip

 78:     test:
 79:       suffix: kok
 80:       requires: kokkos_kernels
 81:       args: -vec_type kokkos

 83: TEST*/