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