1 static char help[] = "Example of using graph partitioning to partition a graph\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 r, N = 10, start, end, *vweights;
11 PetscBool set_vweights = PETSC_FALSE, use_edge_weights = PETSC_FALSE;
12 PetscMPIInt rank;
13 MPI_Comm comm;
14
15 PetscFunctionBeginUser;
16 PetscCall(PetscInitialize(&argc, &args, NULL, help));
17 comm = PETSC_COMM_WORLD;
18 PetscCall(PetscOptionsGetInt(NULL, NULL, "-N", &N, NULL));
19 PetscCallMPI(MPI_Comm_rank(comm, &rank));
20 PetscCall(MatCreate(comm, &A));
21 PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, N, N));
22 PetscCall(MatSetFromOptions(A));
23 PetscCall(MatSeqAIJSetPreallocation(A, 3, NULL));
24 PetscCall(MatMPIAIJSetPreallocation(A, 3, NULL, 2, NULL));
25 PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_vertex_weights", &set_vweights, NULL));
26 PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_use_edge_weights", &use_edge_weights, NULL));
27 /* Create a linear mesh */
28 PetscCall(MatGetOwnershipRange(A, &start, &end));
29 if (set_vweights) {
30 PetscCall(PetscMalloc1(end - start, &vweights));
31 for (r = start; r < end; ++r) vweights[r - start] = rank + 1;
32 }
33 for (r = start; r < end; ++r) {
34 if (r == 0) {
35 PetscInt cols[2];
36 PetscScalar vals[2];
37
38 cols[0] = r;
39 cols[1] = r + 1;
40 vals[0] = 1.0;
41 vals[1] = use_edge_weights ? 2.0 : 1.0;
42
43 PetscCall(MatSetValues(A, 1, &r, 2, cols, vals, INSERT_VALUES));
44 } else if (r == N - 1) {
45 PetscInt cols[2];
46 PetscScalar vals[2];
47
48 cols[0] = r - 1;
49 cols[1] = r;
50 vals[0] = use_edge_weights ? 3.0 : 1.0;
51 vals[1] = 1.0;
52
53 PetscCall(MatSetValues(A, 1, &r, 2, cols, vals, INSERT_VALUES));
54 } else {
55 PetscInt cols[3];
56 PetscScalar vals[3];
57
58 cols[0] = r - 1;
59 cols[1] = r;
60 cols[2] = r + 1;
61 /* ADJ matrix needs to be symmetric */
62 vals[0] = use_edge_weights ? (cols[0] == 0 ? 2.0 : 5.0) : 1.0;
63 vals[1] = 1.0;
64 vals[2] = use_edge_weights ? (cols[2] == N - 1 ? 3.0 : 5.0) : 1.0;
65
66 PetscCall(MatSetValues(A, 1, &r, 3, cols, vals, INSERT_VALUES));
67 }
68 }
69 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
70 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
71
72 PetscCall(MatPartitioningCreate(comm, &part));
73 PetscCall(MatPartitioningSetAdjacency(part, A));
74 if (set_vweights) PetscCall(MatPartitioningSetVertexWeights(part, vweights));
75 if (use_edge_weights) {
76 PetscCall(MatPartitioningSetUseEdgeWeights(part, use_edge_weights));
77
78 PetscCall(MatPartitioningGetUseEdgeWeights(part, &use_edge_weights));
79 PetscCheck(use_edge_weights, comm, PETSC_ERR_ARG_INCOMP, "use_edge_weights flag does not setup correctly ");
80 }
81 PetscCall(MatPartitioningSetFromOptions(part));
82 PetscCall(MatPartitioningApply(part, &is));
83 PetscCall(ISView(is, PETSC_VIEWER_STDOUT_WORLD));
84 PetscCall(ISDestroy(&is));
85 PetscCall(MatPartitioningDestroy(&part));
86
87 PetscCall(MatDestroy(&A));
88 PetscCall(PetscFinalize());
89 return 0;
90 }
91
92 /*TEST
93
94 test:
95 nsize: 3
96 requires: parmetis
97 args: -mat_partitioning_type parmetis
98
99 test:
100 suffix: 2
101 nsize: 3
102 requires: ptscotch
103 args: -mat_partitioning_type ptscotch
104
105 test:
106 suffix: 3
107 nsize: 4
108 requires: party
109 args: -mat_partitioning_type party
110
111 test:
112 suffix: 4
113 nsize: 3
114 requires: chaco
115 args: -mat_partitioning_type chaco
116
117 test:
118 suffix: 5
119 nsize: 3
120 requires: parmetis
121 args: -mat_partitioning_type hierarch -mat_partitioning_hierarchical_nfineparts 3 -mat_partitioning_nparts 10 -N 100
122
123 test:
124 suffix: 6
125 nsize: 3
126 requires: parmetis
127 args: -mat_partitioning_type hierarch -mat_partitioning_hierarchical_nfineparts 3 -mat_partitioning_nparts 10 -N 100 -test_vertex_weights 1 -mat_partitioning_use_edge_weights 1
128
129 test:
130 suffix: 7
131 nsize: 2
132 requires: parmetis
133 args: -mat_partitioning_type hierarch -mat_partitioning_hierarchical_nfineparts 2 -mat_partitioning_nparts 10 -mat_partitioning_hierarchical_fineparttype hierarch -malloc_dump -N 100 -mat_partitioning_improve 1
134
135 test:
136 suffix: 8
137 nsize: 2
138 requires: parmetis
139 args: -mat_partitioning_type parmetis -mat_partitioning_nparts 3 -test_use_edge_weights 1
140
141 test:
142 suffix: 9
143 nsize: 2
144 requires: ptscotch
145 args: -mat_partitioning_type ptscotch -mat_partitioning_nparts 3 -test_use_edge_weights 1 -mat_partitioning_ptscotch_proc_weight 0
146
147 TEST*/
148