Actual source code: partsimple.c

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

  4: typedef struct {
  5:   PetscBool useGrid;        /* Flag to use a grid layout */
  6:   PetscInt  gridDim;        /* The grid dimension */
  7:   PetscInt  nodeGrid[3];    /* Dimension of node grid */
  8:   PetscInt  processGrid[3]; /* Dimension of local process grid on each node */
  9: } PetscPartitioner_Simple;

 11: static PetscErrorCode PetscPartitionerDestroy_Simple(PetscPartitioner part)
 12: {

 16:   PetscFree(part->data);
 17:   return(0);
 18: }

 20: static PetscErrorCode PetscPartitionerView_Simple_ASCII(PetscPartitioner part, PetscViewer viewer)
 21: {
 23:   return(0);
 24: }

 26: static PetscErrorCode PetscPartitionerView_Simple(PetscPartitioner part, PetscViewer viewer)
 27: {
 28:   PetscBool      iascii;

 34:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
 35:   if (iascii) {PetscPartitionerView_Simple_ASCII(part, viewer);}
 36:   return(0);
 37: }

 39: static PetscErrorCode PetscPartitionerSetFromOptions_Simple(PetscOptionItems *PetscOptionsObject, PetscPartitioner part)
 40: {
 41:   PetscPartitioner_Simple *p = (PetscPartitioner_Simple *) part->data;
 42:   PetscInt                 num, i;
 43:   PetscBool                flg;
 44:   PetscErrorCode           ierr;

 47:   for (i = 0; i < 3; ++i) p->processGrid[i] = p->nodeGrid[i] = 1;
 48:   PetscOptionsHead(PetscOptionsObject, "PetscPartitioner Simple Options");
 49:   num  = 3;
 50:   PetscOptionsIntArray("-petscpartitioner_simple_node_grid", "Number of nodes in each dimension", "", p->nodeGrid, &num, &flg);
 51:   if (flg) {p->useGrid = PETSC_TRUE; p->gridDim = num;}
 52:   num  = 3;
 53:   PetscOptionsIntArray("-petscpartitioner_simple_process_grid", "Number of local processes in each dimension for a given node", "", p->processGrid, &num, &flg);
 54:   if (flg) {
 55:     p->useGrid = PETSC_TRUE;
 56:     if (p->gridDim < 0) p->gridDim = num;
 57:     else if (p->gridDim != num) SETERRQ2(PetscObjectComm((PetscObject) part), PETSC_ERR_ARG_INCOMP, "Process grid dimension %D != %D node grid dimension", num, p->gridDim);
 58:   }
 59:   PetscOptionsTail();
 60:   return(0);
 61: }

 63: static PetscErrorCode PetscPartitionerPartition_Simple_Grid(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
 64: {
 65:   PetscPartitioner_Simple *p = (PetscPartitioner_Simple *) part->data;
 66:   const PetscInt          *nodes = p->nodeGrid;
 67:   const PetscInt          *procs = p->processGrid;
 68:   PetscInt                *cellproc, *offsets, cells[3] = {1, 1, 1}, pcells[3] = {1, 1, 1};
 69:   PetscInt                 Np    = 1, Nlc = 1, Nr, np, nk, nj, ni, pk, pj, pi, ck, cj, ci, i;
 70:   MPI_Comm                 comm;
 71:   PetscMPIInt              size;
 72:   PetscErrorCode           ierr;

 75:   if (vertSection)   {PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores vertex weights when using grid partition\n");}
 76:   if (targetSection) {PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores partition weights when using grid partition\n");}
 77:   PetscObjectGetComm((PetscObject) part, &comm);
 78:   MPI_Comm_size(comm, &size);
 79:   /* Check grid */
 80:   for (i = 0; i < 3; ++i) Np *= nodes[i]*procs[i];
 81:   if (nparts != Np)   SETERRQ2(comm, PETSC_ERR_ARG_INCOMP, "Number of partitions %D != %D grid size", nparts, Np);
 82:   if (nparts != size) SETERRQ2(comm, PETSC_ERR_ARG_INCOMP, "Number of partitions %D != %D processes", nparts, size);
 83:   if (numVertices % nparts) SETERRQ2(comm, PETSC_ERR_ARG_INCOMP, "Number of cells %D is not divisible by number of partitions %D", numVertices, nparts);
 84:   for (i = 0; i < p->gridDim; ++i) cells[i] = nodes[i]*procs[i];
 85:   Nr = numVertices / nparts;
 86:   while (Nr > 1) {
 87:     for (i = 0; i < p->gridDim; ++i) {
 88:       cells[i] *= 2;
 89:       Nr       /= 2;
 90:     }
 91:   }
 92:   if (numVertices && Nr != 1) SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Odd number of cells %D. Must be nprocs*2^k", numVertices);
 93:   for (i = 0; i < p->gridDim; ++i) {
 94:     if (cells[i] %  (nodes[i]*procs[i])) SETERRQ3(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "dir %D. Number of cells (%D) mod number of processors %D", i, cells[i], nodes[i]*procs[i]);
 95:     pcells[i] = cells[i] / (nodes[i]*procs[i]); Nlc *= pcells[i];
 96:   }
 97:   /* Compute sizes */
 98:   for (np = 0; np < nparts; ++np) {PetscSectionSetDof(partSection, np, numVertices/nparts);}
 99:   PetscSectionSetUp(partSection);
100:   PetscCalloc1(nparts, &offsets);
101:   for (np = 0; np < nparts; ++np) {PetscSectionGetOffset(partSection, np, &offsets[np]);}
102:   if (!numVertices) pcells[0] = pcells[1] = pcells[2] = 0;
103:   /* Compute partition */
104:   PetscMalloc1(numVertices, &cellproc);
105:   for (nk = 0; nk < nodes[2]; ++nk) {
106:     for (nj = 0; nj < nodes[1]; ++nj) {
107:       for (ni = 0; ni < nodes[0]; ++ni) {
108:         const PetscInt nid = (nk*nodes[1] + nj)*nodes[0] + ni;

110:         for (pk = 0; pk < procs[2]; ++pk) {
111:           for (pj = 0; pj < procs[1]; ++pj) {
112:             for (pi = 0; pi < procs[0]; ++pi) {
113:               const PetscInt pid = ((nid*procs[2] + pk)*procs[1] + pj)*procs[0] + pi;

115:               /* Assume that cells are originally numbered lexicographically */
116:               for (ck = 0; ck < pcells[2]; ++ck) {
117:                 for (cj = 0; cj < pcells[1]; ++cj) {
118:                   for (ci = 0; ci < pcells[0]; ++ci) {
119:                     const PetscInt cid = (((nk*procs[2] + pk)*pcells[2] + ck)*cells[1] + ((nj*procs[1] + pj)*pcells[1] + cj))*cells[0] + (ni*procs[0] + pi)*pcells[0] + ci;

121:                     cellproc[offsets[pid]++] = cid;
122:                   }
123:                 }
124:               }
125:             }
126:           }
127:         }
128:       }
129:     }
130:   }
131:   for (np = 1; np < nparts; ++np) if (offsets[np] - offsets[np-1] != numVertices/nparts) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Offset %D != %D partition size", offsets[np], numVertices/nparts);
132:   PetscFree(offsets);
133:   ISCreateGeneral(PETSC_COMM_SELF, numVertices, cellproc, PETSC_OWN_POINTER, partition);
134:   return(0);
135: }

137: static PetscErrorCode PetscPartitionerPartition_Simple(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
138: {
139:   PetscPartitioner_Simple *p = (PetscPartitioner_Simple *) part->data;
140:   MPI_Comm       comm;
141:   PetscInt       np, *tpwgts = NULL, sumw = 0, numVerticesGlobal  = 0;
142:   PetscMPIInt    size;

146:   if (p->useGrid) {
147:     PetscPartitionerPartition_Simple_Grid(part, nparts, numVertices, start, adjacency, vertSection, targetSection, partSection, partition);
148:     return(0);
149:   }
150:   if (vertSection) {PetscInfo(part,"PETSCPARTITIONERSIMPLE ignores vertex weights\n");}
151:   PetscObjectGetComm((PetscObject) part, &comm);
152:   MPI_Comm_size(comm, &size);
153:   if (targetSection) {
154:     MPIU_Allreduce(&numVertices, &numVerticesGlobal, 1, MPIU_INT, MPI_SUM, comm);
155:     PetscCalloc1(nparts,&tpwgts);
156:     for (np = 0; np < nparts; ++np) {
157:       PetscSectionGetDof(targetSection,np,&tpwgts[np]);
158:       sumw += tpwgts[np];
159:     }
160:     if (sumw) {
161:       PetscInt m,mp;
162:       for (np = 0; np < nparts; ++np) tpwgts[np] = (tpwgts[np]*numVerticesGlobal)/sumw;
163:       for (np = 0, m = -1, mp = 0, sumw = 0; np < nparts; ++np) {
164:         if (m < tpwgts[np]) { m = tpwgts[np]; mp = np; }
165:         sumw += tpwgts[np];
166:       }
167:       if (sumw != numVerticesGlobal) tpwgts[mp] += numVerticesGlobal - sumw;
168:     }
169:     if (!sumw) {PetscFree(tpwgts);}
170:   }

172:   ISCreateStride(PETSC_COMM_SELF, numVertices, 0, 1, partition);
173:   if (size == 1) {
174:     if (tpwgts) {
175:       for (np = 0; np < nparts; ++np) {
176:         PetscSectionSetDof(partSection, np, tpwgts[np]);
177:       }
178:     } else {
179:       for (np = 0; np < nparts; ++np) {
180:         PetscSectionSetDof(partSection, np, numVertices/nparts + ((numVertices % nparts) > np));
181:       }
182:     }
183:   } else {
184:     if (tpwgts) {
185:       Vec         v;
186:       PetscScalar *array;
187:       PetscInt    st,j;
188:       PetscMPIInt rank;

190:       VecCreate(comm,&v);
191:       VecSetSizes(v,numVertices,numVerticesGlobal);
192:       VecSetType(v,VECSTANDARD);
193:       MPI_Comm_rank(comm,&rank);
194:       for (np = 0,st = 0; np < nparts; ++np) {
195:         if (rank == np || (rank == size-1 && size < nparts && np >= size)) {
196:           for (j = 0; j < tpwgts[np]; j++) {
197:             VecSetValue(v,st+j,np,INSERT_VALUES);
198:           }
199:         }
200:         st += tpwgts[np];
201:       }
202:       VecAssemblyBegin(v);
203:       VecAssemblyEnd(v);
204:       VecGetArray(v,&array);
205:       for (j = 0; j < numVertices; ++j) {
206:         PetscSectionAddDof(partSection,PetscRealPart(array[j]),1);
207:       }
208:       VecRestoreArray(v,&array);
209:       VecDestroy(&v);
210:     } else {
211:       PetscMPIInt rank;
212:       PetscInt nvGlobal, *offsets, myFirst, myLast;

214:       PetscMalloc1(size+1,&offsets);
215:       offsets[0] = 0;
216:       MPI_Allgather(&numVertices,1,MPIU_INT,&offsets[1],1,MPIU_INT,comm);
217:       for (np = 2; np <= size; np++) {
218:         offsets[np] += offsets[np-1];
219:       }
220:       nvGlobal = offsets[size];
221:       MPI_Comm_rank(comm,&rank);
222:       myFirst = offsets[rank];
223:       myLast  = offsets[rank + 1] - 1;
224:       PetscFree(offsets);
225:       if (numVertices) {
226:         PetscInt firstPart = 0, firstLargePart = 0;
227:         PetscInt lastPart = 0, lastLargePart = 0;
228:         PetscInt rem = nvGlobal % nparts;
229:         PetscInt pSmall = nvGlobal/nparts;
230:         PetscInt pBig = nvGlobal/nparts + 1;

232:         if (rem) {
233:           firstLargePart = myFirst / pBig;
234:           lastLargePart  = myLast  / pBig;

236:           if (firstLargePart < rem) {
237:             firstPart = firstLargePart;
238:           } else {
239:             firstPart = rem + (myFirst - (rem * pBig)) / pSmall;
240:           }
241:           if (lastLargePart < rem) {
242:             lastPart = lastLargePart;
243:           } else {
244:             lastPart = rem + (myLast - (rem * pBig)) / pSmall;
245:           }
246:         } else {
247:           firstPart = myFirst / (nvGlobal/nparts);
248:           lastPart  = myLast  / (nvGlobal/nparts);
249:         }

251:         for (np = firstPart; np <= lastPart; np++) {
252:           PetscInt PartStart =  np    * (nvGlobal/nparts) + PetscMin(nvGlobal % nparts,np);
253:           PetscInt PartEnd   = (np+1) * (nvGlobal/nparts) + PetscMin(nvGlobal % nparts,np+1);

255:           PartStart = PetscMax(PartStart,myFirst);
256:           PartEnd   = PetscMin(PartEnd,myLast+1);
257:           PetscSectionSetDof(partSection,np,PartEnd-PartStart);
258:         }
259:       }
260:     }
261:   }
262:   PetscFree(tpwgts);
263:   return(0);
264: }

266: static PetscErrorCode PetscPartitionerInitialize_Simple(PetscPartitioner part)
267: {
269:   part->noGraph             = PETSC_TRUE;
270:   part->ops->view           = PetscPartitionerView_Simple;
271:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_Simple;
272:   part->ops->destroy        = PetscPartitionerDestroy_Simple;
273:   part->ops->partition      = PetscPartitionerPartition_Simple;
274:   return(0);
275: }

277: /*MC
278:   PETSCPARTITIONERSIMPLE = "simple" - A PetscPartitioner object

280:   Level: intermediate

282: .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
283: M*/

285: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Simple(PetscPartitioner part)
286: {
287:   PetscPartitioner_Simple *p;
288:   PetscErrorCode           ierr;

292:   PetscNewLog(part, &p);
293:   p->gridDim = -1;
294:   part->data = p;

296:   PetscPartitionerInitialize_Simple(part);
297:   return(0);
298: }