xref: /petsc/src/vec/is/sf/tests/ex7.c (revision 58d68138c660dfb4e9f5b03334792cd4f2ffd7cc)
1 static char help[] = "Test vecscatter of different block sizes across processes\n\n";
2 
3 #include <petscvec.h>
4 int main(int argc, char **argv) {
5   PetscInt           i, bs, n, low, high;
6   PetscMPIInt        nproc, rank;
7   Vec                x, y, z;
8   IS                 ix, iy;
9   VecScatter         vscat;
10   const PetscScalar *yv;
11 
12   PetscFunctionBeginUser;
13   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
14   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &nproc));
15   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
16   PetscCheck(nproc == 2, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This test can only run on two MPI ranks");
17 
18   /* Create an MPI vector x of size 12 on two processes, and set x = {0, 1, 2, .., 11} */
19   PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 6, PETSC_DECIDE, &x));
20   PetscCall(VecGetOwnershipRange(x, &low, &high));
21   for (i = low; i < high; i++) PetscCall(VecSetValue(x, i, (PetscScalar)i, INSERT_VALUES));
22   PetscCall(VecAssemblyBegin(x));
23   PetscCall(VecAssemblyEnd(x));
24 
25   /* Create a seq vector y, and a parallel to sequential (PtoS) vecscatter to scatter x to y */
26   if (rank == 0) {
27     /* On rank 0, seq y is of size 6. We will scatter x[0,1,2,6,7,8] to y[0,1,2,3,4,5] using IS with bs=3 */
28     PetscInt idx[2] = {0, 2};
29     PetscInt idy[2] = {0, 1};
30     n               = 6;
31     bs              = 3;
32     PetscCall(VecCreateSeq(PETSC_COMM_SELF, n, &y));
33     PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, 2, idx, PETSC_COPY_VALUES, &ix));
34     PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, 2, idy, PETSC_COPY_VALUES, &iy));
35   } else {
36     /* On rank 1, seq y is of size 4. We will scatter x[4,5,10,11] to y[0,1,2,3] using IS with bs=2 */
37     PetscInt idx[2] = {2, 5};
38     PetscInt idy[2] = {0, 1};
39     n               = 4;
40     bs              = 2;
41     PetscCall(VecCreateSeq(PETSC_COMM_SELF, n, &y));
42     PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, 2, idx, PETSC_COPY_VALUES, &ix));
43     PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, 2, idy, PETSC_COPY_VALUES, &iy));
44   }
45   PetscCall(VecScatterCreate(x, ix, y, iy, &vscat));
46 
47   /* Do the vecscatter */
48   PetscCall(VecScatterBegin(vscat, x, y, INSERT_VALUES, SCATTER_FORWARD));
49   PetscCall(VecScatterEnd(vscat, x, y, INSERT_VALUES, SCATTER_FORWARD));
50 
51   /* Print y. Since y is sequential, we put y in a parallel z to print its value on both ranks */
52   PetscCall(VecGetArrayRead(y, &yv));
53   PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, n, PETSC_DECIDE, yv, &z));
54   PetscCall(VecView(z, PETSC_VIEWER_STDOUT_WORLD));
55   PetscCall(VecRestoreArrayRead(y, &yv));
56 
57   PetscCall(ISDestroy(&ix));
58   PetscCall(ISDestroy(&iy));
59   PetscCall(VecDestroy(&x));
60   PetscCall(VecDestroy(&y));
61   PetscCall(VecDestroy(&z));
62   PetscCall(VecScatterDestroy(&vscat));
63 
64   PetscCall(PetscFinalize());
65   return 0;
66 }
67 
68 /*TEST
69 
70    test:
71       nsize: 2
72       args:
73       requires:
74 TEST*/
75