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