1 static const char help[] = "Test PetscSFFetchAndOp \n\n"; 2 3 #include <petscvec.h> 4 #include <petscsf.h> 5 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; 17 18 PetscFunctionBeginUser; 19 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 20 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)); 25 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)); 31 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 34 35 PetscCall(VecSet(x, 1.0)); 36 PetscCall(VecSet(y, 2.0)); 37 38 PetscCall(VecGetArrayReadAndMemType(x, &xarray, &xmtype)); 39 PetscCall(VecGetArrayAndMemType(y, &yarray, &ymtype)); 40 PetscCall(VecGetArrayWriteAndMemType(z, &zarray, &zmtype)); 41 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)); 44 45 PetscCall(VecRestoreArrayReadAndMemType(x, &xarray)); 46 PetscCall(VecRestoreArrayAndMemType(y, &yarray)); 47 PetscCall(VecRestoreArrayWriteAndMemType(z, &zarray)); 48 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 } 59 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" 67 68 test: 69 suffix: cuda 70 requires: cuda 71 args: -vec_type cuda 72 73 test: 74 suffix: hip 75 requires: hip 76 args: -vec_type hip 77 78 test: 79 suffix: kok 80 requires: kokkos_kernels 81 args: -vec_type kokkos 82 83 TEST*/ 84