xref: /petsc/src/dm/partitioner/impls/matpart/partmatpart.c (revision c3b5f7ba6bc5ce25a01a67bb37ba5d34b02bbbd7)
1 #include <petscmat.h>
2 #include <petsc/private/partitionerimpl.h>   /*I      "petscpartitioner.h"   I*/
3 
4 typedef struct {
5   MatPartitioning mp;
6 } PetscPartitioner_MatPartitioning;
7 
8 static PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning(PetscPartitioner part, MatPartitioning *mp)
9 {
10   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
11 
12   PetscFunctionBegin;
13   *mp = p->mp;
14   PetscFunctionReturn(0);
15 }
16 
17 /*@C
18   PetscPartitionerMatPartitioningGetMatPartitioning - Get a MatPartitioning instance wrapped by this PetscPartitioner.
19 
20   Not Collective
21 
22   Input Parameters:
23 . part     - The PetscPartitioner
24 
25   Output Parameters:
26 . mp       - The MatPartitioning
27 
28   Level: developer
29 
30 .seealso DMPlexDistribute(), PetscPartitionerCreate()
31 @*/
32 PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp)
33 {
34   PetscFunctionBegin;
35   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
36   PetscValidPointer(mp, 2);
37   PetscUseMethod(part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",(PetscPartitioner,MatPartitioning*),(part,mp));
38   PetscFunctionReturn(0);
39 }
40 
41 static PetscErrorCode PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part)
42 {
43   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
44 
45   PetscFunctionBegin;
46   PetscCall(MatPartitioningDestroy(&p->mp));
47   PetscCall(PetscFree(part->data));
48   PetscFunctionReturn(0);
49 }
50 
51 static PetscErrorCode PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part, PetscViewer viewer)
52 {
53   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
54   PetscViewerFormat                 format;
55 
56   PetscFunctionBegin;
57   PetscCall(PetscViewerGetFormat(viewer, &format));
58   PetscCall(PetscViewerASCIIPrintf(viewer, "MatPartitioning Graph Partitioner:\n"));
59   PetscCall(PetscViewerASCIIPushTab(viewer));
60   if (p->mp) PetscCall(MatPartitioningView(p->mp, viewer));
61   PetscCall(PetscViewerASCIIPopTab(viewer));
62   PetscFunctionReturn(0);
63 }
64 
65 static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer)
66 {
67   PetscBool      iascii;
68 
69   PetscFunctionBegin;
70   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
71   PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
72   PetscCall(PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii));
73   if (iascii) PetscCall(PetscPartitionerView_MatPartitioning_ASCII(part, viewer));
74   PetscFunctionReturn(0);
75 }
76 
77 static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscOptionItems *PetscOptionsObject, PetscPartitioner part)
78 {
79   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
80 
81   PetscFunctionBegin;
82   PetscCall(PetscObjectSetOptionsPrefix((PetscObject)p->mp,((PetscObject)part)->prefix));
83   PetscCall(MatPartitioningSetFromOptions(p->mp));
84   PetscFunctionReturn(0);
85 }
86 
87 static PetscErrorCode PetscPartitionerPartition_MatPartitioning(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *is)
88 {
89   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
90   Mat                               matadj;
91   IS                                is1, is2, is3;
92   PetscReal                         *tpwgts = NULL;
93   PetscInt                          numVerticesGlobal, numEdges;
94   PetscInt                          *i, *j, *vwgt = NULL;
95   MPI_Comm                          comm;
96 
97   PetscFunctionBegin;
98   PetscCall(PetscObjectGetComm((PetscObject)part, &comm));
99 
100   /* TODO: MatCreateMPIAdj should maybe take global number of ROWS */
101   /* TODO: And vertex distribution in PetscPartitionerPartition_ParMetis should be done using PetscSplitOwnership */
102   numVerticesGlobal = PETSC_DECIDE;
103   PetscCall(PetscSplitOwnership(comm, &numVertices, &numVerticesGlobal));
104 
105   /* copy arrays to avoid memory errors because MatMPIAdjSetPreallocation copies just pointers */
106   numEdges = start[numVertices];
107   PetscCall(PetscMalloc1(numVertices+1, &i));
108   PetscCall(PetscMalloc1(numEdges, &j));
109   PetscCall(PetscArraycpy(i, start, numVertices+1));
110   PetscCall(PetscArraycpy(j, adjacency, numEdges));
111 
112   /* construct the adjacency matrix */
113   PetscCall(MatCreateMPIAdj(comm, numVertices, numVerticesGlobal, i, j, NULL, &matadj));
114   PetscCall(MatPartitioningSetAdjacency(p->mp, matadj));
115   PetscCall(MatPartitioningSetNParts(p->mp, nparts));
116 
117   /* calculate partition weights */
118   if (targetSection) {
119     PetscReal sumt;
120     PetscInt  p;
121 
122     sumt = 0.0;
123     PetscCall(PetscMalloc1(nparts,&tpwgts));
124     for (p = 0; p < nparts; ++p) {
125       PetscInt tpd;
126 
127       PetscCall(PetscSectionGetDof(targetSection,p,&tpd));
128       sumt += tpd;
129       tpwgts[p] = tpd;
130     }
131     if (sumt) { /* METIS/ParMETIS do not like exactly zero weight */
132       for (p = 0, sumt = 0.0; p < nparts; ++p) {
133         tpwgts[p] = PetscMax(tpwgts[p],PETSC_SMALL);
134         sumt += tpwgts[p];
135       }
136       for (p = 0; p < nparts; ++p) tpwgts[p] /= sumt;
137       for (p = 0, sumt = 0.0; p < nparts-1; ++p) sumt += tpwgts[p];
138       tpwgts[nparts - 1] = 1. - sumt;
139     } else {
140       PetscCall(PetscFree(tpwgts));
141     }
142   }
143   PetscCall(MatPartitioningSetPartitionWeights(p->mp, tpwgts));
144 
145   /* calculate vertex weights */
146   if (vertSection) {
147     PetscInt v;
148 
149     PetscCall(PetscMalloc1(numVertices,&vwgt));
150     for (v = 0; v < numVertices; ++v) {
151       PetscCall(PetscSectionGetDof(vertSection, v, &vwgt[v]));
152     }
153   }
154   PetscCall(MatPartitioningSetVertexWeights(p->mp, vwgt));
155 
156   /* apply the partitioning */
157   PetscCall(MatPartitioningApply(p->mp, &is1));
158 
159   /* construct the PetscSection */
160   {
161     PetscInt v;
162     const PetscInt *assignment_arr;
163 
164     PetscCall(ISGetIndices(is1, &assignment_arr));
165     for (v = 0; v < numVertices; ++v) PetscCall(PetscSectionAddDof(partSection, assignment_arr[v], 1));
166     PetscCall(ISRestoreIndices(is1, &assignment_arr));
167   }
168 
169   /* convert assignment IS to global numbering IS */
170   PetscCall(ISPartitioningToNumbering(is1, &is2));
171   PetscCall(ISDestroy(&is1));
172 
173   /* renumber IS into local numbering */
174   PetscCall(ISOnComm(is2, PETSC_COMM_SELF, PETSC_USE_POINTER, &is1));
175   PetscCall(ISRenumber(is1, NULL, NULL, &is3));
176   PetscCall(ISDestroy(&is1));
177   PetscCall(ISDestroy(&is2));
178 
179   /* invert IS */
180   PetscCall(ISSetPermutation(is3));
181   PetscCall(ISInvertPermutation(is3, numVertices, &is1));
182   PetscCall(ISDestroy(&is3));
183 
184   PetscCall(MatDestroy(&matadj));
185   *is = is1;
186   PetscFunctionReturn(0);
187 }
188 
189 static PetscErrorCode PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part)
190 {
191   PetscFunctionBegin;
192   part->ops->view           = PetscPartitionerView_MatPartitioning;
193   part->ops->setfromoptions = PetscPartitionerSetFromOptions_MatPartitioning;
194   part->ops->destroy        = PetscPartitionerDestroy_MatPartitioning;
195   part->ops->partition      = PetscPartitionerPartition_MatPartitioning;
196   PetscCall(PetscObjectComposeFunction((PetscObject)part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning));
197   PetscFunctionReturn(0);
198 }
199 
200 /*MC
201   PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object
202 
203   Level: developer
204 
205 .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
206 M*/
207 
208 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)
209 {
210   PetscPartitioner_MatPartitioning  *p;
211 
212   PetscFunctionBegin;
213   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
214   PetscCall(PetscNewLog(part, &p));
215   part->data = p;
216   PetscCall(PetscPartitionerInitialize_MatPartitioning(part));
217   PetscCall(MatPartitioningCreate(PetscObjectComm((PetscObject)part), &p->mp));
218   PetscFunctionReturn(0);
219 }
220