xref: /petsc/src/mat/tutorials/ex17.c (revision 6a5217c03994f2d95bb2e6dbd8bed42381aeb015)
1 static char help[] = "Example of using graph partitioning with a matrix in which some procs have empty ownership\n\n";
2 
3 #include <petscmat.h>
4 
5 int main(int argc, char **args)
6 {
7   Mat             A;
8   MatPartitioning part;
9   IS              is;
10   PetscInt        i,m,N,rstart,rend,nemptyranks,*emptyranks,nbigranks,*bigranks;
11   PetscMPIInt     rank,size;
12   PetscErrorCode  ierr;
13 
14   PetscCall(PetscInitialize(&argc,&args,(char*)0,help));
15   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&rank));
16   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
17 
18   nemptyranks = 10;
19   nbigranks   = 10;
20   PetscCall(PetscMalloc2(nemptyranks,&emptyranks,nbigranks,&bigranks));
21 
22   ierr = PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"Partitioning example options",NULL);PetscCall(ierr);
23   PetscCall(PetscOptionsIntArray("-emptyranks","Ranks to be skipped by partition","",emptyranks,&nemptyranks,NULL));
24   PetscCall(PetscOptionsIntArray("-bigranks","Ranks to be overloaded","",bigranks,&nbigranks,NULL));
25   ierr = PetscOptionsEnd();PetscCall(ierr);
26 
27   m = 1;
28   for (i=0; i<nemptyranks; i++) if (rank == emptyranks[i]) m = 0;
29   for (i=0; i<nbigranks; i++) if (rank == bigranks[i]) m = 5;
30 
31   PetscCall(MatCreate(PETSC_COMM_WORLD,&A));
32   PetscCall(MatSetSizes(A,m,m,PETSC_DECIDE,PETSC_DECIDE));
33   PetscCall(MatSetFromOptions(A));
34   PetscCall(MatSeqAIJSetPreallocation(A,3,NULL));
35   PetscCall(MatMPIAIJSetPreallocation(A,3,NULL,2,NULL));
36   PetscCall(MatSeqBAIJSetPreallocation(A,1,3,NULL));
37   PetscCall(MatMPIBAIJSetPreallocation(A,1,3,NULL,2,NULL));
38   PetscCall(MatSeqSBAIJSetPreallocation(A,1,2,NULL));
39   PetscCall(MatMPISBAIJSetPreallocation(A,1,2,NULL,1,NULL));
40 
41   PetscCall(MatGetSize(A,NULL,&N));
42   PetscCall(MatGetOwnershipRange(A,&rstart,&rend));
43   for (i=rstart; i<rend; i++) {
44     const PetscInt    cols[] = {(i+N-1)%N,i,(i+1)%N};
45     const PetscScalar vals[] = {1,1,1};
46     PetscCall(MatSetValues(A,1,&i,3,cols,vals,INSERT_VALUES));
47   }
48   PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
49   PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
50   PetscCall(MatView(A,PETSC_VIEWER_STDOUT_WORLD));
51 
52   PetscCall(MatPartitioningCreate(PETSC_COMM_WORLD,&part));
53   PetscCall(MatPartitioningSetAdjacency(part,A));
54   PetscCall(MatPartitioningSetFromOptions(part));
55   PetscCall(MatPartitioningApply(part,&is));
56   PetscCall(ISView(is,PETSC_VIEWER_STDOUT_WORLD));
57   PetscCall(ISDestroy(&is));
58   PetscCall(MatPartitioningDestroy(&part));
59   PetscCall(MatDestroy(&A));
60   PetscCall(PetscFree2(emptyranks,bigranks));
61   PetscCall(PetscFinalize());
62   return 0;
63 }
64 
65 /*TEST
66 
67    test:
68       nsize: 8
69       args: -emptyranks 0,2,4 -bigranks 1,3,7 -mat_partitioning_type average
70       # cannot test with external package partitioners since they produce different results on different systems
71 
72 TEST*/
73