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