Actual source code: ex65.c

  1: const char help[] = "Test VecGetLocalVector()";

  3: #include <petscvec.h>

  5: int main(int argc, char **argv)
  6: {
  7:   Vec                global, global_copy, local;
  8:   PetscMPIInt        rank;
  9:   PetscMemType       memtype;
 10:   PetscScalar       *array;
 11:   PetscInt           N = 10;
 12:   const PetscScalar *copy_array;

 14:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 15:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
 16:   PetscCall(VecCreate(PETSC_COMM_WORLD, &global));
 17:   PetscCall(PetscObjectSetName((PetscObject)global, "global"));
 18:   PetscCall(VecSetType(global, VECMPICUDA));
 19:   PetscCall(VecSetSizes(global, rank == 0 ? N : 0, N));
 20:   PetscCall(VecSetRandom(global, NULL));
 21:   PetscCall(VecDuplicate(global, &global_copy));
 22:   PetscCall(VecCopy(global, global_copy));

 24:   PetscCall(VecGetArrayRead(global_copy, &copy_array));
 25:   PetscCall(VecGetArrayAndMemType(global, &array, &memtype));
 26:   PetscCall(VecRestoreArrayAndMemType(global, &array));
 27:   if (rank == 0) {
 28:     PetscOffloadMask mask;

 30:     PetscCall(VecGetOffloadMask(global, &mask));
 31:     PetscCheck(mask == PETSC_OFFLOAD_GPU, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Unexpected offload state");
 32:   }

 34:   PetscCall(VecCreateLocalVector(global, &local));
 35:   PetscCall(PetscObjectSetName((PetscObject)local, "local"));
 36:   PetscCall(VecGetLocalVector(global, local));
 37:   if (rank == 0) {
 38:     const PetscScalar *local_array;
 39:     PetscOffloadMask   mask;

 41:     PetscCall(VecGetOffloadMask(local, &mask));
 42:     PetscCheck(mask == PETSC_OFFLOAD_GPU, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Local vector has synced with host");
 43:     PetscCall(VecGetOffloadMask(global, &mask));
 44:     PetscCheck(mask == PETSC_OFFLOAD_GPU, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Global vector has synced with host");

 46:     PetscCall(VecGetArrayRead(local, &local_array));
 47:     for (PetscInt i = 0; i < N; i++) {
 48:       PetscCheck(copy_array[i] == local_array[i], PETSC_COMM_SELF, PETSC_ERR_PLIB, "VecGetLocalVector() value mismatch: local[%" PetscInt_FMT "] = %g, global[%" PetscInt_FMT "] = %g", i, (double)PetscRealPart(local_array[i]), i, PetscRealPart(copy_array[i]));
 49:     }
 50:     PetscCall(VecRestoreArrayRead(local, &local_array));
 51:   }
 52:   PetscCall(VecRestoreLocalVector(global, local));
 53:   PetscCall(VecDestroy(&local));
 54:   PetscCall(VecRestoreArrayRead(global_copy, &copy_array));
 55:   PetscCall(VecDestroy(&global_copy));
 56:   PetscCall(VecDestroy(&global));
 57:   PetscCall(PetscFinalize());
 58:   return 0;
 59: }

 61: /*TEST

 63:   test:
 64:     requires: cuda
 65:     nsize: {{1 2}}
 66:     suffix: 0

 68: TEST*/