xref: /petsc/src/dm/partitioner/impls/simple/partsimple.c (revision efbe7e8a80d07327753dbe0b33efee01e046af3f)
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       PetscInt m,mp;
56       for (np = 0; np < nparts; ++np) tpwgts[np] = (tpwgts[np]*numVerticesGlobal)/sumw;
57       for (np = 0, m = -1, mp = 0, sumw = 0; np < nparts; ++np) {
58         if (m < tpwgts[np]) { m = tpwgts[np]; mp = np; }
59         sumw += tpwgts[np];
60       }
61       if (sumw != numVerticesGlobal) tpwgts[mp] += numVerticesGlobal - sumw;
62     }
63     if (!sumw) {ierr = PetscFree(tpwgts);CHKERRQ(ierr);}
64   }
65 
66   ierr = ISCreateStride(PETSC_COMM_SELF, numVertices, 0, 1, partition);CHKERRQ(ierr);
67   if (size == 1) {
68     if (tpwgts) {
69       for (np = 0; np < nparts; ++np) {
70         ierr = PetscSectionSetDof(partSection, np, tpwgts[np]);CHKERRQ(ierr);
71       }
72     } else {
73       for (np = 0; np < nparts; ++np) {
74         ierr = PetscSectionSetDof(partSection, np, numVertices/nparts + ((numVertices % nparts) > np));CHKERRQ(ierr);
75       }
76     }
77   } else {
78     if (tpwgts) {
79       Vec         v;
80       PetscScalar *array;
81       PetscInt    st,j;
82       PetscMPIInt rank;
83 
84       ierr = VecCreate(comm,&v);CHKERRQ(ierr);
85       ierr = VecSetSizes(v,numVertices,numVerticesGlobal);CHKERRQ(ierr);
86       ierr = VecSetType(v,VECSTANDARD);CHKERRQ(ierr);
87       ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
88       for (np = 0,st = 0; np < nparts; ++np) {
89         if (rank == np || (rank == size-1 && size < nparts && np >= size)) {
90           for (j = 0; j < tpwgts[np]; j++) {
91             ierr = VecSetValue(v,st+j,np,INSERT_VALUES);CHKERRQ(ierr);
92           }
93         }
94         st += tpwgts[np];
95       }
96       ierr = VecAssemblyBegin(v);CHKERRQ(ierr);
97       ierr = VecAssemblyEnd(v);CHKERRQ(ierr);
98       ierr = VecGetArray(v,&array);CHKERRQ(ierr);
99       for (j = 0; j < numVertices; ++j) {
100         ierr = PetscSectionAddDof(partSection,PetscRealPart(array[j]),1);CHKERRQ(ierr);
101       }
102       ierr = VecRestoreArray(v,&array);CHKERRQ(ierr);
103       ierr = VecDestroy(&v);CHKERRQ(ierr);
104     } else {
105       PetscMPIInt rank;
106       PetscInt nvGlobal, *offsets, myFirst, myLast;
107 
108       ierr = PetscMalloc1(size+1,&offsets);CHKERRQ(ierr);
109       offsets[0] = 0;
110       ierr = MPI_Allgather(&numVertices,1,MPIU_INT,&offsets[1],1,MPIU_INT,comm);CHKERRQ(ierr);
111       for (np = 2; np <= size; np++) {
112         offsets[np] += offsets[np-1];
113       }
114       nvGlobal = offsets[size];
115       ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
116       myFirst = offsets[rank];
117       myLast  = offsets[rank + 1] - 1;
118       ierr = PetscFree(offsets);CHKERRQ(ierr);
119       if (numVertices) {
120         PetscInt firstPart = 0, firstLargePart = 0;
121         PetscInt lastPart = 0, lastLargePart = 0;
122         PetscInt rem = nvGlobal % nparts;
123         PetscInt pSmall = nvGlobal/nparts;
124         PetscInt pBig = nvGlobal/nparts + 1;
125 
126         if (rem) {
127           firstLargePart = myFirst / pBig;
128           lastLargePart  = myLast  / pBig;
129 
130           if (firstLargePart < rem) {
131             firstPart = firstLargePart;
132           } else {
133             firstPart = rem + (myFirst - (rem * pBig)) / pSmall;
134           }
135           if (lastLargePart < rem) {
136             lastPart = lastLargePart;
137           } else {
138             lastPart = rem + (myLast - (rem * pBig)) / pSmall;
139           }
140         } else {
141           firstPart = myFirst / (nvGlobal/nparts);
142           lastPart  = myLast  / (nvGlobal/nparts);
143         }
144 
145         for (np = firstPart; np <= lastPart; np++) {
146           PetscInt PartStart =  np    * (nvGlobal/nparts) + PetscMin(nvGlobal % nparts,np);
147           PetscInt PartEnd   = (np+1) * (nvGlobal/nparts) + PetscMin(nvGlobal % nparts,np+1);
148 
149           PartStart = PetscMax(PartStart,myFirst);
150           PartEnd   = PetscMin(PartEnd,myLast+1);
151           ierr = PetscSectionSetDof(partSection,np,PartEnd-PartStart);CHKERRQ(ierr);
152         }
153       }
154     }
155   }
156   ierr = PetscFree(tpwgts);CHKERRQ(ierr);
157   PetscFunctionReturn(0);
158 }
159 
160 static PetscErrorCode PetscPartitionerInitialize_Simple(PetscPartitioner part)
161 {
162   PetscFunctionBegin;
163   part->noGraph        = PETSC_TRUE;
164   part->ops->view      = PetscPartitionerView_Simple;
165   part->ops->destroy   = PetscPartitionerDestroy_Simple;
166   part->ops->partition = PetscPartitionerPartition_Simple;
167   PetscFunctionReturn(0);
168 }
169 
170 /*MC
171   PETSCPARTITIONERSIMPLE = "simple" - A PetscPartitioner object
172 
173   Level: intermediate
174 
175 .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
176 M*/
177 
178 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Simple(PetscPartitioner part)
179 {
180   PetscPartitioner_Simple *p;
181   PetscErrorCode           ierr;
182 
183   PetscFunctionBegin;
184   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
185   ierr       = PetscNewLog(part, &p);CHKERRQ(ierr);
186   part->data = p;
187 
188   ierr = PetscPartitionerInitialize_Simple(part);CHKERRQ(ierr);
189   PetscFunctionReturn(0);
190 }
191