xref: /petsc/src/mat/tutorials/ex17.c (revision d5b43468fb8780a8feea140ccd6fa3e6a50411cc)
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 
13   PetscFunctionBeginUser;
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   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Partitioning example options", NULL);
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   PetscOptionsEnd();
26 
27   m = 1;
28   for (i = 0; i < nemptyranks; i++)
29     if (rank == emptyranks[i]) m = 0;
30   for (i = 0; i < nbigranks; i++)
31     if (rank == bigranks[i]) m = 5;
32 
33   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
34   PetscCall(MatSetSizes(A, m, m, PETSC_DECIDE, PETSC_DECIDE));
35   PetscCall(MatSetFromOptions(A));
36   PetscCall(MatSeqAIJSetPreallocation(A, 3, NULL));
37   PetscCall(MatMPIAIJSetPreallocation(A, 3, NULL, 2, NULL));
38   PetscCall(MatSeqBAIJSetPreallocation(A, 1, 3, NULL));
39   PetscCall(MatMPIBAIJSetPreallocation(A, 1, 3, NULL, 2, NULL));
40   PetscCall(MatSeqSBAIJSetPreallocation(A, 1, 2, NULL));
41   PetscCall(MatMPISBAIJSetPreallocation(A, 1, 2, NULL, 1, NULL));
42 
43   PetscCall(MatGetSize(A, NULL, &N));
44   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
45   for (i = rstart; i < rend; i++) {
46     const PetscInt    cols[] = {(i + N - 1) % N, i, (i + 1) % N};
47     const PetscScalar vals[] = {1, 1, 1};
48     PetscCall(MatSetValues(A, 1, &i, 3, cols, vals, INSERT_VALUES));
49   }
50   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
51   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
52   PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
53 
54   PetscCall(MatPartitioningCreate(PETSC_COMM_WORLD, &part));
55   PetscCall(MatPartitioningSetAdjacency(part, A));
56   PetscCall(MatPartitioningSetFromOptions(part));
57   PetscCall(MatPartitioningApply(part, &is));
58   PetscCall(ISView(is, PETSC_VIEWER_STDOUT_WORLD));
59   PetscCall(ISDestroy(&is));
60   PetscCall(MatPartitioningDestroy(&part));
61   PetscCall(MatDestroy(&A));
62   PetscCall(PetscFree2(emptyranks, bigranks));
63   PetscCall(PetscFinalize());
64   return 0;
65 }
66 
67 /*TEST
68 
69    test:
70       nsize: 8
71       args: -emptyranks 0,2,4 -bigranks 1,3,7 -mat_partitioning_type average
72       # cannot test with external package partitioners since they produce different results on different systems
73 
74 TEST*/
75