xref: /petsc/src/dm/partitioner/impls/shell/partshell.c (revision bcda9346efad4e5ba2d553af84eb238771ba1e25)
1 #include <petsc/private/partitionerimpl.h> /*I "petscpartitioner.h" I*/
2 
3 typedef struct {
4   PetscSection section;   /* Sizes for each partition */
5   IS           partition; /* Points in each partition */
6   PetscBool    random;    /* Flag for a random partition */
7 } PetscPartitioner_Shell;
8 
PetscPartitionerReset_Shell(PetscPartitioner part)9 static PetscErrorCode PetscPartitionerReset_Shell(PetscPartitioner part)
10 {
11   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
12 
13   PetscFunctionBegin;
14   PetscCall(PetscSectionDestroy(&p->section));
15   PetscCall(ISDestroy(&p->partition));
16   PetscFunctionReturn(PETSC_SUCCESS);
17 }
18 
PetscPartitionerDestroy_Shell(PetscPartitioner part)19 static PetscErrorCode PetscPartitionerDestroy_Shell(PetscPartitioner part)
20 {
21   PetscFunctionBegin;
22   PetscCall(PetscPartitionerReset_Shell(part));
23   PetscCall(PetscFree(part->data));
24   PetscFunctionReturn(PETSC_SUCCESS);
25 }
26 
PetscPartitionerView_Shell_ASCII(PetscPartitioner part,PetscViewer viewer)27 static PetscErrorCode PetscPartitionerView_Shell_ASCII(PetscPartitioner part, PetscViewer viewer)
28 {
29   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
30 
31   PetscFunctionBegin;
32   if (p->random) {
33     PetscCall(PetscViewerASCIIPushTab(viewer));
34     PetscCall(PetscViewerASCIIPrintf(viewer, "using random partition\n"));
35     PetscCall(PetscViewerASCIIPopTab(viewer));
36   }
37   PetscFunctionReturn(PETSC_SUCCESS);
38 }
39 
PetscPartitionerView_Shell(PetscPartitioner part,PetscViewer viewer)40 static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer)
41 {
42   PetscBool isascii;
43 
44   PetscFunctionBegin;
45   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
46   PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
47   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
48   if (isascii) PetscCall(PetscPartitionerView_Shell_ASCII(part, viewer));
49   PetscFunctionReturn(PETSC_SUCCESS);
50 }
51 
PetscPartitionerSetFromOptions_Shell(PetscPartitioner part,PetscOptionItems PetscOptionsObject)52 static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscPartitioner part, PetscOptionItems PetscOptionsObject)
53 {
54   PetscInt    sizes[16], points[1024];
55   PetscInt    Npart = 16, Npoints = 1024;
56   PetscBool   random = PETSC_FALSE, set, flgSizes, flgPoints;
57   PetscMPIInt rank;
58 
59   PetscFunctionBegin;
60   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank));
61   PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Shell Options");
62   PetscCall(PetscPartitionerShellGetRandom(part, &random));
63   PetscCall(PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set));
64   if (set) PetscCall(PetscPartitionerShellSetRandom(part, random));
65   PetscCall(PetscOptionsIntArray("-petscpartitioner_shell_sizes", "The size of each partition on rank 0", "PetscPartitionerShellSetPartition", sizes, &Npart, &flgSizes));
66   PetscCall(PetscOptionsIntArray("-petscpartitioner_shell_points", "The points in each partition on rank 0", "PetscPartitionerShellSetPartition", points, &Npoints, &flgPoints));
67   PetscCheck(!(flgSizes ^ flgPoints), PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Must specify both the partition sizes and points");
68   if (flgSizes) {
69     PetscInt Np = 0;
70 
71     for (PetscInt i = 0; i < Npart; ++i) Np += sizes[i];
72     PetscCheck(Np == Npoints, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Number of input points %" PetscInt_FMT " != %" PetscInt_FMT " sum of partition sizes", Npoints, Np);
73     if (!rank) PetscCall(PetscPartitionerShellSetPartition(part, Npart, sizes, points));
74     else {
75       PetscCall(PetscArrayzero(sizes, Npart));
76       PetscCall(PetscPartitionerShellSetPartition(part, Npart, sizes, points));
77     }
78   }
79   PetscOptionsHeadEnd();
80   PetscFunctionReturn(PETSC_SUCCESS);
81 }
82 
PetscPartitionerPartition_Shell(PetscPartitioner part,PetscInt nparts,PetscInt numVertices,PetscInt start[],PetscInt adjacency[],PetscSection vertSection,PetscSection edgeSection,PetscSection targetSection,PetscSection partSection,IS * partition)83 static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *partition)
84 {
85   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
86   PetscInt                np;
87 
88   PetscFunctionBegin;
89   if (p->random) {
90     PetscRandom r;
91     PetscInt   *sizes, *points, v, p;
92     PetscMPIInt rank;
93 
94     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank));
95     PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &r));
96     PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)nparts));
97     PetscCall(PetscRandomSetFromOptions(r));
98     PetscCall(PetscCalloc2(nparts, &sizes, numVertices, &points));
99     for (v = 0; v < numVertices; ++v) points[v] = v;
100     for (p = 0; p < nparts; ++p) sizes[p] = numVertices / nparts + (PetscInt)(p < numVertices % nparts);
101     for (v = numVertices - 1; v > 0; --v) {
102       PetscReal val;
103       PetscInt  w, tmp;
104 
105       PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)(v + 1)));
106       PetscCall(PetscRandomGetValueReal(r, &val));
107       w         = PetscFloorReal(val);
108       tmp       = points[v];
109       points[v] = points[w];
110       points[w] = tmp;
111     }
112     PetscCall(PetscRandomDestroy(&r));
113     PetscCall(PetscPartitionerShellSetPartition(part, nparts, sizes, points));
114     PetscCall(PetscFree2(sizes, points));
115   }
116   PetscCheck(p->section, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Shell partitioner information not provided. Please call PetscPartitionerShellSetPartition()");
117   PetscCall(PetscSectionGetChart(p->section, NULL, &np));
118   PetscCheck(nparts == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of requested partitions %" PetscInt_FMT " != configured partitions %" PetscInt_FMT, nparts, np);
119   PetscCall(ISGetLocalSize(p->partition, &np));
120   PetscCheck(numVertices == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of input vertices %" PetscInt_FMT " != configured vertices %" PetscInt_FMT, numVertices, np);
121   PetscCall(PetscSectionCopy(p->section, partSection));
122   *partition = p->partition;
123   PetscCall(PetscObjectReference((PetscObject)p->partition));
124   PetscFunctionReturn(PETSC_SUCCESS);
125 }
126 
PetscPartitionerInitialize_Shell(PetscPartitioner part)127 static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part)
128 {
129   PetscFunctionBegin;
130   part->noGraph             = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */
131   part->ops->view           = PetscPartitionerView_Shell;
132   part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell;
133   part->ops->reset          = PetscPartitionerReset_Shell;
134   part->ops->destroy        = PetscPartitionerDestroy_Shell;
135   part->ops->partition      = PetscPartitionerPartition_Shell;
136   PetscFunctionReturn(PETSC_SUCCESS);
137 }
138 
139 /*MC
140   PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object
141 
142   Level: intermediate
143 
144   Options Database Keys:
145 .  -petscpartitioner_shell_random - Use a random partition
146 
147 .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
148 M*/
149 
PetscPartitionerCreate_Shell(PetscPartitioner part)150 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part)
151 {
152   PetscPartitioner_Shell *p;
153 
154   PetscFunctionBegin;
155   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
156   PetscCall(PetscNew(&p));
157   part->data = p;
158 
159   PetscCall(PetscPartitionerInitialize_Shell(part));
160   p->random = PETSC_FALSE;
161   PetscFunctionReturn(PETSC_SUCCESS);
162 }
163 
164 /*@
165   PetscPartitionerShellSetPartition - Set an artificial partition for a mesh
166 
167   Collective
168 
169   Input Parameters:
170 + part   - The `PetscPartitioner`
171 . size   - The number of partitions
172 . sizes  - array of length size (or `NULL`) providing the number of points in each partition
173 - points - array of length sum(sizes) (may be `NULL` iff sizes is `NULL`), a permutation of the points that groups those assigned to each partition in order (i.e., partition 0 first, partition 1 next, etc.)
174 
175   Level: developer
176 
177   Note:
178   It is safe to free the sizes and points arrays after use in this routine.
179 
180 .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()`
181 @*/
PetscPartitionerShellSetPartition(PetscPartitioner part,PetscInt size,const PetscInt sizes[],const PetscInt points[])182 PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[])
183 {
184   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
185   PetscInt                proc, numPoints;
186 
187   PetscFunctionBegin;
188   PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL);
189   if (sizes) PetscAssertPointer(sizes, 3);
190   if (points) PetscAssertPointer(points, 4);
191   PetscCall(PetscSectionDestroy(&p->section));
192   PetscCall(ISDestroy(&p->partition));
193   PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)part), &p->section));
194   PetscCall(PetscSectionSetChart(p->section, 0, size));
195   if (sizes) {
196     for (proc = 0; proc < size; ++proc) PetscCall(PetscSectionSetDof(p->section, proc, sizes[proc]));
197   }
198   PetscCall(PetscSectionSetUp(p->section));
199   PetscCall(PetscSectionGetStorageSize(p->section, &numPoints));
200   PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)part), numPoints, points, PETSC_COPY_VALUES, &p->partition));
201   PetscFunctionReturn(PETSC_SUCCESS);
202 }
203 
204 /*@
205   PetscPartitionerShellSetRandom - Set the flag to use a random partition
206 
207   Collective
208 
209   Input Parameters:
210 + part   - The `PetscPartitioner`
211 - random - The flag to use a random partition
212 
213   Level: intermediate
214 
215 .seealso: `PetscPartitionerShellGetRandom()`, `PetscPartitionerCreate()`
216 @*/
PetscPartitionerShellSetRandom(PetscPartitioner part,PetscBool random)217 PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random)
218 {
219   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
220 
221   PetscFunctionBegin;
222   PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL);
223   p->random = random;
224   PetscFunctionReturn(PETSC_SUCCESS);
225 }
226 
227 /*@
228   PetscPartitionerShellGetRandom - get the flag to use a random partition
229 
230   Collective
231 
232   Input Parameter:
233 . part - The `PetscPartitioner`
234 
235   Output Parameter:
236 . random - The flag to use a random partition
237 
238   Level: intermediate
239 
240 .seealso: `PetscPartitionerShellSetRandom()`, `PetscPartitionerCreate()`
241 @*/
PetscPartitionerShellGetRandom(PetscPartitioner part,PetscBool * random)242 PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random)
243 {
244   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
245 
246   PetscFunctionBegin;
247   PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL);
248   PetscAssertPointer(random, 2);
249   *random = p->random;
250   PetscFunctionReturn(PETSC_SUCCESS);
251 }
252