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
main(int argc,char ** args)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, NULL, 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