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: {
13: PetscFunctionBegin;
14: PetscCall(PetscFree(part->data));
15: PetscFunctionReturn(PETSC_SUCCESS);
16: }
18: static PetscErrorCode PetscPartitionerView_Simple_ASCII(PetscPartitioner part, PetscViewer viewer)
19: {
20: PetscFunctionBegin;
21: PetscFunctionReturn(PETSC_SUCCESS);
22: }
24: static PetscErrorCode PetscPartitionerView_Simple(PetscPartitioner part, PetscViewer viewer)
25: {
26: PetscBool iascii;
28: PetscFunctionBegin;
31: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
32: if (iascii) PetscCall(PetscPartitionerView_Simple_ASCII(part, viewer));
33: PetscFunctionReturn(PETSC_SUCCESS);
34: }
36: static PetscErrorCode PetscPartitionerSetFromOptions_Simple(PetscPartitioner part, PetscOptionItems *PetscOptionsObject)
37: {
38: PetscPartitioner_Simple *p = (PetscPartitioner_Simple *)part->data;
39: PetscInt num, i;
40: PetscBool flg;
42: PetscFunctionBegin;
43: for (i = 0; i < 3; ++i) p->processGrid[i] = p->nodeGrid[i] = 1;
44: PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Simple Options");
45: num = 3;
46: PetscCall(PetscOptionsIntArray("-petscpartitioner_simple_node_grid", "Number of nodes in each dimension", "", p->nodeGrid, &num, &flg));
47: if (flg) {
48: p->useGrid = PETSC_TRUE;
49: p->gridDim = num;
50: }
51: num = 3;
52: PetscCall(PetscOptionsIntArray("-petscpartitioner_simple_process_grid", "Number of local processes in each dimension for a given node", "", p->processGrid, &num, &flg));
53: if (flg) {
54: p->useGrid = PETSC_TRUE;
55: if (p->gridDim < 0) p->gridDim = num;
56: else PetscCheck(p->gridDim == num, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_INCOMP, "Process grid dimension %" PetscInt_FMT " != %" PetscInt_FMT " node grid dimension", num, p->gridDim);
57: }
58: PetscOptionsHeadEnd();
59: PetscFunctionReturn(PETSC_SUCCESS);
60: }
62: static PetscErrorCode PetscPartitionerPartition_Simple_Grid(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
63: {
64: PetscPartitioner_Simple *p = (PetscPartitioner_Simple *)part->data;
65: const PetscInt *nodes = p->nodeGrid;
66: const PetscInt *procs = p->processGrid;
67: PetscInt *cellproc, *offsets, cells[3] = {1, 1, 1}, pcells[3] = {1, 1, 1};
68: PetscInt Np = 1, Nr, np, nk, nj, ni, pk, pj, pi, ck, cj, ci, i;
69: MPI_Comm comm;
70: PetscMPIInt size;
72: PetscFunctionBegin;
73: if (vertSection) PetscCall(PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores vertex weights when using grid partition\n"));
74: if (targetSection) PetscCall(PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores partition weights when using grid partition\n"));
75: PetscCall(PetscObjectGetComm((PetscObject)part, &comm));
76: PetscCallMPI(MPI_Comm_size(comm, &size));
77: /* Check grid */
78: for (i = 0; i < 3; ++i) Np *= nodes[i] * procs[i];
79: PetscCheck(nparts == Np, comm, PETSC_ERR_ARG_INCOMP, "Number of partitions %" PetscInt_FMT " != %" PetscInt_FMT " grid size", nparts, Np);
80: PetscCheck(nparts == size, comm, PETSC_ERR_ARG_INCOMP, "Number of partitions %" PetscInt_FMT " != %d processes", nparts, size);
81: PetscCheck(numVertices % nparts == 0, comm, PETSC_ERR_ARG_INCOMP, "Number of cells %" PetscInt_FMT " is not divisible by number of partitions %" PetscInt_FMT, numVertices, nparts);
82: for (i = 0; i < p->gridDim; ++i) cells[i] = nodes[i] * procs[i];
83: Nr = numVertices / nparts;
84: while (Nr > 1) {
85: for (i = 0; i < p->gridDim; ++i) {
86: cells[i] *= 2;
87: Nr /= 2;
88: }
89: }
90: PetscCheck(!numVertices || Nr == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Odd number of cells %" PetscInt_FMT ". Must be nprocs*2^k", numVertices);
91: for (i = 0; i < p->gridDim; ++i) {
92: PetscCheck(cells[i] % (nodes[i] * procs[i]) == 0, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "dir %" PetscInt_FMT ". Number of cells (%" PetscInt_FMT ") mod number of processors %" PetscInt_FMT, i, cells[i], nodes[i] * procs[i]);
93: pcells[i] = cells[i] / (nodes[i] * procs[i]);
94: }
95: /* Compute sizes */
96: for (np = 0; np < nparts; ++np) PetscCall(PetscSectionSetDof(partSection, np, numVertices / nparts));
97: PetscCall(PetscSectionSetUp(partSection));
98: PetscCall(PetscCalloc1(nparts, &offsets));
99: for (np = 0; np < nparts; ++np) PetscCall(PetscSectionGetOffset(partSection, np, &offsets[np]));
100: if (!numVertices) pcells[0] = pcells[1] = pcells[2] = 0;
101: /* Compute partition */
102: PetscCall(PetscMalloc1(numVertices, &cellproc));
103: for (nk = 0; nk < nodes[2]; ++nk) {
104: for (nj = 0; nj < nodes[1]; ++nj) {
105: for (ni = 0; ni < nodes[0]; ++ni) {
106: const PetscInt nid = (nk * nodes[1] + nj) * nodes[0] + ni;
108: for (pk = 0; pk < procs[2]; ++pk) {
109: for (pj = 0; pj < procs[1]; ++pj) {
110: for (pi = 0; pi < procs[0]; ++pi) {
111: const PetscInt pid = ((nid * procs[2] + pk) * procs[1] + pj) * procs[0] + pi;
113: /* Assume that cells are originally numbered lexicographically */
114: for (ck = 0; ck < pcells[2]; ++ck) {
115: for (cj = 0; cj < pcells[1]; ++cj) {
116: for (ci = 0; ci < pcells[0]; ++ci) {
117: 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;
119: cellproc[offsets[pid]++] = cid;
120: }
121: }
122: }
123: }
124: }
125: }
126: }
127: }
128: }
129: for (np = 1; np < nparts; ++np) PetscCheck(offsets[np] - offsets[np - 1] == numVertices / nparts, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Offset %" PetscInt_FMT " != %" PetscInt_FMT " partition size", offsets[np], numVertices / nparts);
130: PetscCall(PetscFree(offsets));
131: PetscCall(ISCreateGeneral(PETSC_COMM_SELF, numVertices, cellproc, PETSC_OWN_POINTER, partition));
132: PetscFunctionReturn(PETSC_SUCCESS);
133: }
135: static PetscErrorCode PetscPartitionerPartition_Simple(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *partition)
136: {
137: PetscPartitioner_Simple *p = (PetscPartitioner_Simple *)part->data;
138: MPI_Comm comm;
139: PetscInt np, *tpwgts = NULL, sumw = 0, numVerticesGlobal = 0;
140: PetscMPIInt size;
142: PetscFunctionBegin;
143: if (p->useGrid) {
144: PetscCall(PetscPartitionerPartition_Simple_Grid(part, nparts, numVertices, start, adjacency, vertSection, targetSection, partSection, partition));
145: PetscFunctionReturn(PETSC_SUCCESS);
146: }
147: if (vertSection) PetscCall(PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores vertex weights\n"));
148: PetscCall(PetscObjectGetComm((PetscObject)part, &comm));
149: PetscCallMPI(MPI_Comm_size(comm, &size));
150: if (targetSection) {
151: PetscCallMPI(MPIU_Allreduce(&numVertices, &numVerticesGlobal, 1, MPIU_INT, MPI_SUM, comm));
152: PetscCall(PetscCalloc1(nparts, &tpwgts));
153: for (np = 0; np < nparts; ++np) {
154: PetscCall(PetscSectionGetDof(targetSection, np, &tpwgts[np]));
155: sumw += tpwgts[np];
156: }
157: if (sumw) {
158: PetscInt m, mp;
159: for (np = 0; np < nparts; ++np) tpwgts[np] = (tpwgts[np] * numVerticesGlobal) / sumw;
160: for (np = 0, m = -1, mp = 0, sumw = 0; np < nparts; ++np) {
161: if (m < tpwgts[np]) {
162: m = tpwgts[np];
163: mp = np;
164: }
165: sumw += tpwgts[np];
166: }
167: if (sumw != numVerticesGlobal) tpwgts[mp] += numVerticesGlobal - sumw;
168: }
169: if (!sumw) PetscCall(PetscFree(tpwgts));
170: }
172: PetscCall(ISCreateStride(PETSC_COMM_SELF, numVertices, 0, 1, partition));
173: if (size == 1) {
174: if (tpwgts) {
175: for (np = 0; np < nparts; ++np) PetscCall(PetscSectionSetDof(partSection, np, tpwgts[np]));
176: } else {
177: for (np = 0; np < nparts; ++np) PetscCall(PetscSectionSetDof(partSection, np, numVertices / nparts + ((numVertices % nparts) > np)));
178: }
179: } else {
180: if (tpwgts) {
181: Vec v;
182: PetscScalar *array;
183: PetscInt st, j;
184: PetscMPIInt rank;
186: PetscCall(VecCreate(comm, &v));
187: PetscCall(VecSetSizes(v, numVertices, numVerticesGlobal));
188: PetscCall(VecSetType(v, VECSTANDARD));
189: PetscCallMPI(MPI_Comm_rank(comm, &rank));
190: for (np = 0, st = 0; np < nparts; ++np) {
191: if (rank == np || (rank == size - 1 && size < nparts && np >= size)) {
192: for (j = 0; j < tpwgts[np]; j++) PetscCall(VecSetValue(v, st + j, np, INSERT_VALUES));
193: }
194: st += tpwgts[np];
195: }
196: PetscCall(VecAssemblyBegin(v));
197: PetscCall(VecAssemblyEnd(v));
198: PetscCall(VecGetArray(v, &array));
199: for (j = 0; j < numVertices; ++j) PetscCall(PetscSectionAddDof(partSection, PetscRealPart(array[j]), 1));
200: PetscCall(VecRestoreArray(v, &array));
201: PetscCall(VecDestroy(&v));
202: } else {
203: PetscMPIInt rank;
204: PetscInt nvGlobal, *offsets, myFirst, myLast;
206: PetscCall(PetscMalloc1(size + 1, &offsets));
207: offsets[0] = 0;
208: PetscCallMPI(MPI_Allgather(&numVertices, 1, MPIU_INT, &offsets[1], 1, MPIU_INT, comm));
209: for (np = 2; np <= size; np++) offsets[np] += offsets[np - 1];
210: nvGlobal = offsets[size];
211: PetscCallMPI(MPI_Comm_rank(comm, &rank));
212: myFirst = offsets[rank];
213: myLast = offsets[rank + 1] - 1;
214: PetscCall(PetscFree(offsets));
215: if (numVertices) {
216: PetscInt firstPart = 0, firstLargePart = 0;
217: PetscInt lastPart = 0, lastLargePart = 0;
218: PetscInt rem = nvGlobal % nparts;
219: PetscInt pSmall = nvGlobal / nparts;
220: PetscInt pBig = nvGlobal / nparts + 1;
222: if (rem) {
223: firstLargePart = myFirst / pBig;
224: lastLargePart = myLast / pBig;
226: if (firstLargePart < rem) {
227: firstPart = firstLargePart;
228: } else {
229: firstPart = rem + (myFirst - (rem * pBig)) / pSmall;
230: }
231: if (lastLargePart < rem) {
232: lastPart = lastLargePart;
233: } else {
234: lastPart = rem + (myLast - (rem * pBig)) / pSmall;
235: }
236: } else {
237: firstPart = myFirst / (nvGlobal / nparts);
238: lastPart = myLast / (nvGlobal / nparts);
239: }
241: for (np = firstPart; np <= lastPart; np++) {
242: PetscInt PartStart = np * (nvGlobal / nparts) + PetscMin(nvGlobal % nparts, np);
243: PetscInt PartEnd = (np + 1) * (nvGlobal / nparts) + PetscMin(nvGlobal % nparts, np + 1);
245: PartStart = PetscMax(PartStart, myFirst);
246: PartEnd = PetscMin(PartEnd, myLast + 1);
247: PetscCall(PetscSectionSetDof(partSection, np, PartEnd - PartStart));
248: }
249: }
250: }
251: }
252: PetscCall(PetscFree(tpwgts));
253: PetscFunctionReturn(PETSC_SUCCESS);
254: }
256: static PetscErrorCode PetscPartitionerInitialize_Simple(PetscPartitioner part)
257: {
258: PetscFunctionBegin;
259: part->noGraph = PETSC_TRUE;
260: part->ops->view = PetscPartitionerView_Simple;
261: part->ops->setfromoptions = PetscPartitionerSetFromOptions_Simple;
262: part->ops->destroy = PetscPartitionerDestroy_Simple;
263: part->ops->partition = PetscPartitionerPartition_Simple;
264: PetscFunctionReturn(PETSC_SUCCESS);
265: }
267: /*MC
268: PETSCPARTITIONERSIMPLE = "simple" - A PetscPartitioner object
270: Level: intermediate
272: .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
273: M*/
275: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Simple(PetscPartitioner part)
276: {
277: PetscPartitioner_Simple *p;
279: PetscFunctionBegin;
281: PetscCall(PetscNew(&p));
282: p->gridDim = -1;
283: part->data = p;
285: PetscCall(PetscPartitionerInitialize_Simple(part));
286: PetscFunctionReturn(PETSC_SUCCESS);
287: }