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:   PetscSectionDestroy(&p->section);
 14:   ISDestroy(&p->partition);
 15:   return 0;
 16: }

 18: static PetscErrorCode PetscPartitionerDestroy_Shell(PetscPartitioner part)
 19: {
 20:   PetscPartitionerReset_Shell(part);
 21:   PetscFree(part->data);
 22:   return 0;
 23: }

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

 29:   if (p->random) {
 30:     PetscViewerASCIIPushTab(viewer);
 31:     PetscViewerASCIIPrintf(viewer, "using random partition\n");
 32:     PetscViewerASCIIPopTab(viewer);
 33:   }
 34:   return 0;
 35: }

 37: static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer)
 38: {
 39:   PetscBool iascii;

 43:   PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii);
 44:   if (iascii) PetscPartitionerView_Shell_ASCII(part, viewer);
 45:   return 0;
 46: }

 48: static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscPartitioner part, PetscOptionItems *PetscOptionsObject)
 49: {
 50:   PetscBool random = PETSC_FALSE, set;

 52:   PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Shell Options");
 53:   PetscPartitionerShellGetRandom(part, &random);
 54:   PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set);
 55:   if (set) PetscPartitionerShellSetRandom(part, random);
 56:   PetscOptionsHeadEnd();
 57:   return 0;
 58: }

 60: static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
 61: {
 62:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
 63:   PetscInt                np;

 65:   if (p->random) {
 66:     PetscRandom r;
 67:     PetscInt   *sizes, *points, v, p;
 68:     PetscMPIInt rank;

 70:     MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank);
 71:     PetscRandomCreate(PETSC_COMM_SELF, &r);
 72:     PetscRandomSetInterval(r, 0.0, (PetscScalar)nparts);
 73:     PetscRandomSetFromOptions(r);
 74:     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;

 81:       PetscRandomSetInterval(r, 0.0, (PetscScalar)(v + 1));
 82:       PetscRandomGetValueReal(r, &val);
 83:       w         = PetscFloorReal(val);
 84:       tmp       = points[v];
 85:       points[v] = points[w];
 86:       points[w] = tmp;
 87:     }
 88:     PetscRandomDestroy(&r);
 89:     PetscPartitionerShellSetPartition(part, nparts, sizes, points);
 90:     PetscFree2(sizes, points);
 91:   }
 93:   PetscSectionGetChart(p->section, NULL, &np);
 95:   ISGetLocalSize(p->partition, &np);
 97:   PetscSectionCopy(p->section, partSection);
 98:   *partition = p->partition;
 99:   PetscObjectReference((PetscObject)p->partition);
100:   return 0;
101: }

103: static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part)
104: {
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:   return 0;
112: }

114: /*MC
115:   PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object

117:   Level: intermediate

119:   Options Database Keys:
120: .  -petscpartitioner_shell_random - Use a random partition

122: .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
123: M*/

125: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part)
126: {
127:   PetscPartitioner_Shell *p;

130:   PetscNew(&p);
131:   part->data = p;

133:   PetscPartitionerInitialize_Shell(part);
134:   p->random = PETSC_FALSE;
135:   return 0;
136: }

138: /*@C
139:   PetscPartitionerShellSetPartition - Set an artificial partition for a mesh

141:   Collective on PetscPartitioner

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.)

149:   Level: developer

151:   Notes:
152:     It is safe to free the sizes and points arrays after use in this routine.

154: .seealso `DMPlexDistribute()`, `PetscPartitionerCreate()`
155: @*/
156: PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[])
157: {
158:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
159:   PetscInt                proc, numPoints;

164:   PetscSectionDestroy(&p->section);
165:   ISDestroy(&p->partition);
166:   PetscSectionCreate(PetscObjectComm((PetscObject)part), &p->section);
167:   PetscSectionSetChart(p->section, 0, size);
168:   if (sizes) {
169:     for (proc = 0; proc < size; ++proc) PetscSectionSetDof(p->section, proc, sizes[proc]);
170:   }
171:   PetscSectionSetUp(p->section);
172:   PetscSectionGetStorageSize(p->section, &numPoints);
173:   ISCreateGeneral(PetscObjectComm((PetscObject)part), numPoints, points, PETSC_COPY_VALUES, &p->partition);
174:   return 0;
175: }

177: /*@
178:   PetscPartitionerShellSetRandom - Set the flag to use a random partition

180:   Collective on PetscPartitioner

182:   Input Parameters:
183: + part   - The PetscPartitioner
184: - random - The flag to use a random partition

186:   Level: intermediate

188: .seealso `PetscPartitionerShellGetRandom()`, `PetscPartitionerCreate()`
189: @*/
190: PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random)
191: {
192:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

195:   p->random = random;
196:   return 0;
197: }

199: /*@
200:   PetscPartitionerShellGetRandom - get the flag to use a random partition

202:   Collective on PetscPartitioner

204:   Input Parameter:
205: . part   - The PetscPartitioner

207:   Output Parameter:
208: . random - The flag to use a random partition

210:   Level: intermediate

212: .seealso `PetscPartitionerShellSetRandom()`, `PetscPartitionerCreate()`
213: @*/
214: PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random)
215: {
216:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

220:   *random = p->random;
221:   return 0;
222: }