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