Actual source code: partshell.c

  1: #include <petsc/private/partitionerimpl.h>

  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;

  9: static PetscErrorCode PetscPartitionerReset_Shell(PetscPartitioner part)
 10: {
 11:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

 13:   PetscFunctionBegin;
 14:   PetscCall(PetscSectionDestroy(&p->section));
 15:   PetscCall(ISDestroy(&p->partition));
 16:   PetscFunctionReturn(PETSC_SUCCESS);
 17: }

 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: }

 27: static PetscErrorCode PetscPartitionerView_Shell_ASCII(PetscPartitioner part, PetscViewer viewer)
 28: {
 29:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

 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: }

 40: static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer)
 41: {
 42:   PetscBool iascii;

 44:   PetscFunctionBegin;
 47:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
 48:   if (iascii) PetscCall(PetscPartitionerView_Shell_ASCII(part, viewer));
 49:   PetscFunctionReturn(PETSC_SUCCESS);
 50: }

 52: static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscPartitioner part, PetscOptionItems *PetscOptionsObject)
 53: {
 54:   PetscBool random = PETSC_FALSE, set;

 56:   PetscFunctionBegin;
 57:   PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Shell Options");
 58:   PetscCall(PetscPartitionerShellGetRandom(part, &random));
 59:   PetscCall(PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set));
 60:   if (set) PetscCall(PetscPartitionerShellSetRandom(part, random));
 61:   PetscOptionsHeadEnd();
 62:   PetscFunctionReturn(PETSC_SUCCESS);
 63: }

 65: static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *partition)
 66: {
 67:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
 68:   PetscInt                np;

 70:   PetscFunctionBegin;
 71:   if (p->random) {
 72:     PetscRandom r;
 73:     PetscInt   *sizes, *points, v, p;
 74:     PetscMPIInt rank;

 76:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank));
 77:     PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &r));
 78:     PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)nparts));
 79:     PetscCall(PetscRandomSetFromOptions(r));
 80:     PetscCall(PetscCalloc2(nparts, &sizes, numVertices, &points));
 81:     for (v = 0; v < numVertices; ++v) points[v] = v;
 82:     for (p = 0; p < nparts; ++p) sizes[p] = numVertices / nparts + (PetscInt)(p < numVertices % nparts);
 83:     for (v = numVertices - 1; v > 0; --v) {
 84:       PetscReal val;
 85:       PetscInt  w, tmp;

 87:       PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)(v + 1)));
 88:       PetscCall(PetscRandomGetValueReal(r, &val));
 89:       w         = PetscFloorReal(val);
 90:       tmp       = points[v];
 91:       points[v] = points[w];
 92:       points[w] = tmp;
 93:     }
 94:     PetscCall(PetscRandomDestroy(&r));
 95:     PetscCall(PetscPartitionerShellSetPartition(part, nparts, sizes, points));
 96:     PetscCall(PetscFree2(sizes, points));
 97:   }
 98:   PetscCheck(p->section, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Shell partitioner information not provided. Please call PetscPartitionerShellSetPartition()");
 99:   PetscCall(PetscSectionGetChart(p->section, NULL, &np));
100:   PetscCheck(nparts == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of requested partitions %" PetscInt_FMT " != configured partitions %" PetscInt_FMT, nparts, np);
101:   PetscCall(ISGetLocalSize(p->partition, &np));
102:   PetscCheck(numVertices == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of input vertices %" PetscInt_FMT " != configured vertices %" PetscInt_FMT, numVertices, np);
103:   PetscCall(PetscSectionCopy(p->section, partSection));
104:   *partition = p->partition;
105:   PetscCall(PetscObjectReference((PetscObject)p->partition));
106:   PetscFunctionReturn(PETSC_SUCCESS);
107: }

109: static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part)
110: {
111:   PetscFunctionBegin;
112:   part->noGraph             = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */
113:   part->ops->view           = PetscPartitionerView_Shell;
114:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell;
115:   part->ops->reset          = PetscPartitionerReset_Shell;
116:   part->ops->destroy        = PetscPartitionerDestroy_Shell;
117:   part->ops->partition      = PetscPartitionerPartition_Shell;
118:   PetscFunctionReturn(PETSC_SUCCESS);
119: }

121: /*MC
122:   PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object

124:   Level: intermediate

126:   Options Database Keys:
127: .  -petscpartitioner_shell_random - Use a random partition

129: .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
130: M*/

132: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part)
133: {
134:   PetscPartitioner_Shell *p;

136:   PetscFunctionBegin;
138:   PetscCall(PetscNew(&p));
139:   part->data = p;

141:   PetscCall(PetscPartitionerInitialize_Shell(part));
142:   p->random = PETSC_FALSE;
143:   PetscFunctionReturn(PETSC_SUCCESS);
144: }

146: /*@C
147:   PetscPartitionerShellSetPartition - Set an artificial partition for a mesh

149:   Collective

151:   Input Parameters:
152: + part   - The `PetscPartitioner`
153: . size   - The number of partitions
154: . sizes  - array of length size (or `NULL`) providing the number of points in each partition
155: - 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.)

157:   Level: developer

159:   Note:
160:   It is safe to free the sizes and points arrays after use in this routine.

162: .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()`
163: @*/
164: PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[])
165: {
166:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
167:   PetscInt                proc, numPoints;

169:   PetscFunctionBegin;
171:   if (sizes) PetscAssertPointer(sizes, 3);
172:   if (points) PetscAssertPointer(points, 4);
173:   PetscCall(PetscSectionDestroy(&p->section));
174:   PetscCall(ISDestroy(&p->partition));
175:   PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)part), &p->section));
176:   PetscCall(PetscSectionSetChart(p->section, 0, size));
177:   if (sizes) {
178:     for (proc = 0; proc < size; ++proc) PetscCall(PetscSectionSetDof(p->section, proc, sizes[proc]));
179:   }
180:   PetscCall(PetscSectionSetUp(p->section));
181:   PetscCall(PetscSectionGetStorageSize(p->section, &numPoints));
182:   PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)part), numPoints, points, PETSC_COPY_VALUES, &p->partition));
183:   PetscFunctionReturn(PETSC_SUCCESS);
184: }

186: /*@
187:   PetscPartitionerShellSetRandom - Set the flag to use a random partition

189:   Collective

191:   Input Parameters:
192: + part   - The `PetscPartitioner`
193: - random - The flag to use a random partition

195:   Level: intermediate

197: .seealso: `PetscPartitionerShellGetRandom()`, `PetscPartitionerCreate()`
198: @*/
199: PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random)
200: {
201:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

203:   PetscFunctionBegin;
205:   p->random = random;
206:   PetscFunctionReturn(PETSC_SUCCESS);
207: }

209: /*@
210:   PetscPartitionerShellGetRandom - get the flag to use a random partition

212:   Collective

214:   Input Parameter:
215: . part - The `PetscPartitioner`

217:   Output Parameter:
218: . random - The flag to use a random partition

220:   Level: intermediate

222: .seealso: `PetscPartitionerShellSetRandom()`, `PetscPartitionerCreate()`
223: @*/
224: PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random)
225: {
226:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

228:   PetscFunctionBegin;
230:   PetscAssertPointer(random, 2);
231:   *random = p->random;
232:   PetscFunctionReturn(PETSC_SUCCESS);
233: }