xref: /petsc/src/dm/partitioner/impls/simple/partsimple.c (revision bef158480efac06de457f7a665168877ab3c2fd7)
1 #include <petscvec.h>
2 #include <petsc/private/partitionerimpl.h>        /*I "petscpartitioner.h" I*/
3 
4 typedef struct {
5   PetscInt dummy;
6 } PetscPartitioner_Simple;
7 
8 static PetscErrorCode PetscPartitionerDestroy_Simple(PetscPartitioner part)
9 {
10   PetscErrorCode ierr;
11 
12   PetscFunctionBegin;
13   ierr = PetscFree(part->data);CHKERRQ(ierr);
14   PetscFunctionReturn(0);
15 }
16 
17 static PetscErrorCode PetscPartitionerView_Simple_ASCII(PetscPartitioner part, PetscViewer viewer)
18 {
19   PetscFunctionBegin;
20   PetscFunctionReturn(0);
21 }
22 
23 static PetscErrorCode PetscPartitionerView_Simple(PetscPartitioner part, PetscViewer viewer)
24 {
25   PetscBool      iascii;
26   PetscErrorCode ierr;
27 
28   PetscFunctionBegin;
29   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
30   PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
31   ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);CHKERRQ(ierr);
32   if (iascii) {ierr = PetscPartitionerView_Simple_ASCII(part, viewer);CHKERRQ(ierr);}
33   PetscFunctionReturn(0);
34 }
35 
36 static PetscErrorCode PetscPartitionerPartition_Simple(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
37 {
38   MPI_Comm       comm;
39   PetscInt       np, *tpwgts = NULL, sumw = 0, numVerticesGlobal  = 0;
40   PetscMPIInt    size;
41   PetscErrorCode ierr;
42 
43   PetscFunctionBegin;
44   if (vertSection) { ierr = PetscInfo(part,"PETSCPARTITIONERSIMPLE ignores vertex weights\n");CHKERRQ(ierr); }
45   comm = PetscObjectComm((PetscObject)part);
46   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
47   if (targetSection) {
48     ierr = MPIU_Allreduce(&numVertices, &numVerticesGlobal, 1, MPIU_INT, MPI_SUM, comm);CHKERRQ(ierr);
49     ierr = PetscCalloc1(nparts,&tpwgts);CHKERRQ(ierr);
50     for (np = 0; np < nparts; ++np) {
51       ierr = PetscSectionGetDof(targetSection,np,&tpwgts[np]);CHKERRQ(ierr);
52       sumw += tpwgts[np];
53     }
54     if (!sumw) {
55       ierr = PetscFree(tpwgts);CHKERRQ(ierr);
56     } else {
57       PetscInt m,mp;
58       for (np = 0; np < nparts; ++np) tpwgts[np] = (tpwgts[np]*numVerticesGlobal)/sumw;
59       for (np = 0, m = -1, mp = 0, sumw = 0; np < nparts; ++np) {
60         if (m < tpwgts[np]) { m = tpwgts[np]; mp = np; }
61         sumw += tpwgts[np];
62       }
63       if (sumw != numVerticesGlobal) tpwgts[mp] += numVerticesGlobal - sumw;
64     }
65   }
66 
67   ierr = ISCreateStride(PETSC_COMM_SELF, numVertices, 0, 1, partition);CHKERRQ(ierr);
68   if (size == 1) {
69     if (tpwgts) {
70       for (np = 0; np < nparts; ++np) {
71         ierr = PetscSectionSetDof(partSection, np, tpwgts[np]);CHKERRQ(ierr);
72       }
73     } else {
74       for (np = 0; np < nparts; ++np) {
75         ierr = PetscSectionSetDof(partSection, np, numVertices/nparts + ((numVertices % nparts) > np));CHKERRQ(ierr);
76       }
77     }
78   } else {
79     if (tpwgts) {
80       Vec         v;
81       PetscScalar *array;
82       PetscInt    st,j;
83       PetscMPIInt rank;
84 
85       ierr = VecCreate(comm,&v);CHKERRQ(ierr);
86       ierr = VecSetSizes(v,numVertices,numVerticesGlobal);CHKERRQ(ierr);
87       ierr = VecSetType(v,VECSTANDARD);CHKERRQ(ierr);
88       ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
89       for (np = 0,st = 0; np < nparts; ++np) {
90         if (rank == np || (rank == size-1 && size < nparts && np >= size)) {
91           for (j = 0; j < tpwgts[np]; j++) {
92             ierr = VecSetValue(v,st+j,np,INSERT_VALUES);CHKERRQ(ierr);
93           }
94         }
95         st += tpwgts[np];
96       }
97       ierr = VecAssemblyBegin(v);CHKERRQ(ierr);
98       ierr = VecAssemblyEnd(v);CHKERRQ(ierr);
99       ierr = VecGetArray(v,&array);CHKERRQ(ierr);
100       for (j = 0; j < numVertices; ++j) {
101         ierr = PetscSectionAddDof(partSection,PetscRealPart(array[j]),1);CHKERRQ(ierr);
102       }
103       ierr = VecRestoreArray(v,&array);CHKERRQ(ierr);
104       ierr = VecDestroy(&v);CHKERRQ(ierr);
105     } else {
106       PetscMPIInt rank;
107       PetscInt nvGlobal, *offsets, myFirst, myLast;
108 
109       ierr = PetscMalloc1(size+1,&offsets);CHKERRQ(ierr);
110       offsets[0] = 0;
111       ierr = MPI_Allgather(&numVertices,1,MPIU_INT,&offsets[1],1,MPIU_INT,comm);CHKERRQ(ierr);
112       for (np = 2; np <= size; np++) {
113         offsets[np] += offsets[np-1];
114       }
115       nvGlobal = offsets[size];
116       ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
117       myFirst = offsets[rank];
118       myLast  = offsets[rank + 1] - 1;
119       ierr = PetscFree(offsets);CHKERRQ(ierr);
120       if (numVertices) {
121         PetscInt firstPart = 0, firstLargePart = 0;
122         PetscInt lastPart = 0, lastLargePart = 0;
123         PetscInt rem = nvGlobal % nparts;
124         PetscInt pSmall = nvGlobal/nparts;
125         PetscInt pBig = nvGlobal/nparts + 1;
126 
127         if (rem) {
128           firstLargePart = myFirst / pBig;
129           lastLargePart  = myLast  / pBig;
130 
131           if (firstLargePart < rem) {
132             firstPart = firstLargePart;
133           } else {
134             firstPart = rem + (myFirst - (rem * pBig)) / pSmall;
135           }
136           if (lastLargePart < rem) {
137             lastPart = lastLargePart;
138           } else {
139             lastPart = rem + (myLast - (rem * pBig)) / pSmall;
140           }
141         } else {
142           firstPart = myFirst / (nvGlobal/nparts);
143           lastPart  = myLast  / (nvGlobal/nparts);
144         }
145 
146         for (np = firstPart; np <= lastPart; np++) {
147           PetscInt PartStart =  np    * (nvGlobal/nparts) + PetscMin(nvGlobal % nparts,np);
148           PetscInt PartEnd   = (np+1) * (nvGlobal/nparts) + PetscMin(nvGlobal % nparts,np+1);
149 
150           PartStart = PetscMax(PartStart,myFirst);
151           PartEnd   = PetscMin(PartEnd,myLast+1);
152           ierr = PetscSectionSetDof(partSection,np,PartEnd-PartStart);CHKERRQ(ierr);
153         }
154       }
155     }
156   }
157   ierr = PetscFree(tpwgts);CHKERRQ(ierr);
158   PetscFunctionReturn(0);
159 }
160 
161 static PetscErrorCode PetscPartitionerInitialize_Simple(PetscPartitioner part)
162 {
163   PetscFunctionBegin;
164   part->noGraph        = PETSC_TRUE;
165   part->ops->view      = PetscPartitionerView_Simple;
166   part->ops->destroy   = PetscPartitionerDestroy_Simple;
167   part->ops->partition = PetscPartitionerPartition_Simple;
168   PetscFunctionReturn(0);
169 }
170 
171 /*MC
172   PETSCPARTITIONERSIMPLE = "simple" - A PetscPartitioner object
173 
174   Level: intermediate
175 
176 .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
177 M*/
178 
179 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Simple(PetscPartitioner part)
180 {
181   PetscPartitioner_Simple *p;
182   PetscErrorCode           ierr;
183 
184   PetscFunctionBegin;
185   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
186   ierr       = PetscNewLog(part, &p);CHKERRQ(ierr);
187   part->data = p;
188 
189   ierr = PetscPartitionerInitialize_Simple(part);CHKERRQ(ierr);
190   PetscFunctionReturn(0);
191 }
192