1 #include <petscvec.h> 2 #include <petsc/private/partitionerimpl.h> /*I "petscpartitioner.h" I*/ 3 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; 10 11 static PetscErrorCode PetscPartitionerDestroy_Simple(PetscPartitioner part) 12 { 13 PetscErrorCode ierr; 14 15 PetscFunctionBegin; 16 ierr = PetscFree(part->data);CHKERRQ(ierr); 17 PetscFunctionReturn(0); 18 } 19 20 static PetscErrorCode PetscPartitionerView_Simple_ASCII(PetscPartitioner part, PetscViewer viewer) 21 { 22 PetscFunctionBegin; 23 PetscFunctionReturn(0); 24 } 25 26 static PetscErrorCode PetscPartitionerView_Simple(PetscPartitioner part, PetscViewer viewer) 27 { 28 PetscBool iascii; 29 PetscErrorCode ierr; 30 31 PetscFunctionBegin; 32 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 33 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 34 ierr = PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);CHKERRQ(ierr); 35 if (iascii) {ierr = PetscPartitionerView_Simple_ASCII(part, viewer);CHKERRQ(ierr);} 36 PetscFunctionReturn(0); 37 } 38 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; 45 46 PetscFunctionBegin; 47 for (i = 0; i < 3; ++i) p->processGrid[i] = p->nodeGrid[i] = 1; 48 ierr = PetscOptionsHead(PetscOptionsObject, "PetscPartitioner Simple Options");CHKERRQ(ierr); 49 num = 3; 50 ierr = PetscOptionsIntArray("-petscpartitioner_simple_node_grid", "Number of nodes in each dimension", "", p->nodeGrid, &num, &flg);CHKERRQ(ierr); 51 if (flg) {p->useGrid = PETSC_TRUE; p->gridDim = num;} 52 num = 3; 53 ierr = PetscOptionsIntArray("-petscpartitioner_simple_process_grid", "Number of local processes in each dimension for a given node", "", p->processGrid, &num, &flg);CHKERRQ(ierr); 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 ierr = PetscOptionsTail();CHKERRQ(ierr); 60 PetscFunctionReturn(0); 61 } 62 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; 73 74 PetscFunctionBegin; 75 if (vertSection) {ierr = PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores vertex weights when using grid partition\n");CHKERRQ(ierr);} 76 if (targetSection) {ierr = PetscInfo(part, "PETSCPARTITIONERSIMPLE ignores partition weights when using grid partition\n");CHKERRQ(ierr);} 77 ierr = PetscObjectGetComm((PetscObject) part, &comm);CHKERRQ(ierr); 78 ierr = MPI_Comm_size(comm, &size);CHKERRMPI(ierr); 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) {ierr = PetscSectionSetDof(partSection, np, numVertices/nparts);CHKERRQ(ierr);} 99 ierr = PetscSectionSetUp(partSection);CHKERRQ(ierr); 100 ierr = PetscCalloc1(nparts, &offsets);CHKERRQ(ierr); 101 for (np = 0; np < nparts; ++np) {ierr = PetscSectionGetOffset(partSection, np, &offsets[np]);CHKERRQ(ierr);} 102 if (!numVertices) pcells[0] = pcells[1] = pcells[2] = 0; 103 /* Compute partition */ 104 ierr = PetscMalloc1(numVertices, &cellproc);CHKERRQ(ierr); 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; 109 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; 114 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; 120 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 ierr = PetscFree(offsets);CHKERRQ(ierr); 133 ierr = ISCreateGeneral(PETSC_COMM_SELF, numVertices, cellproc, PETSC_OWN_POINTER, partition);CHKERRQ(ierr); 134 PetscFunctionReturn(0); 135 } 136 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; 143 PetscErrorCode ierr; 144 145 PetscFunctionBegin; 146 if (p->useGrid) { 147 ierr = PetscPartitionerPartition_Simple_Grid(part, nparts, numVertices, start, adjacency, vertSection, targetSection, partSection, partition);CHKERRQ(ierr); 148 PetscFunctionReturn(0); 149 } 150 if (vertSection) {ierr = PetscInfo(part,"PETSCPARTITIONERSIMPLE ignores vertex weights\n");CHKERRQ(ierr);} 151 ierr = PetscObjectGetComm((PetscObject) part, &comm);CHKERRQ(ierr); 152 ierr = MPI_Comm_size(comm, &size);CHKERRMPI(ierr); 153 if (targetSection) { 154 ierr = MPIU_Allreduce(&numVertices, &numVerticesGlobal, 1, MPIU_INT, MPI_SUM, comm);CHKERRMPI(ierr); 155 ierr = PetscCalloc1(nparts,&tpwgts);CHKERRQ(ierr); 156 for (np = 0; np < nparts; ++np) { 157 ierr = PetscSectionGetDof(targetSection,np,&tpwgts[np]);CHKERRQ(ierr); 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) {ierr = PetscFree(tpwgts);CHKERRQ(ierr);} 170 } 171 172 ierr = ISCreateStride(PETSC_COMM_SELF, numVertices, 0, 1, partition);CHKERRQ(ierr); 173 if (size == 1) { 174 if (tpwgts) { 175 for (np = 0; np < nparts; ++np) { 176 ierr = PetscSectionSetDof(partSection, np, tpwgts[np]);CHKERRQ(ierr); 177 } 178 } else { 179 for (np = 0; np < nparts; ++np) { 180 ierr = PetscSectionSetDof(partSection, np, numVertices/nparts + ((numVertices % nparts) > np));CHKERRQ(ierr); 181 } 182 } 183 } else { 184 if (tpwgts) { 185 Vec v; 186 PetscScalar *array; 187 PetscInt st,j; 188 PetscMPIInt rank; 189 190 ierr = VecCreate(comm,&v);CHKERRQ(ierr); 191 ierr = VecSetSizes(v,numVertices,numVerticesGlobal);CHKERRQ(ierr); 192 ierr = VecSetType(v,VECSTANDARD);CHKERRQ(ierr); 193 ierr = MPI_Comm_rank(comm,&rank);CHKERRMPI(ierr); 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 ierr = VecSetValue(v,st+j,np,INSERT_VALUES);CHKERRQ(ierr); 198 } 199 } 200 st += tpwgts[np]; 201 } 202 ierr = VecAssemblyBegin(v);CHKERRQ(ierr); 203 ierr = VecAssemblyEnd(v);CHKERRQ(ierr); 204 ierr = VecGetArray(v,&array);CHKERRQ(ierr); 205 for (j = 0; j < numVertices; ++j) { 206 ierr = PetscSectionAddDof(partSection,PetscRealPart(array[j]),1);CHKERRQ(ierr); 207 } 208 ierr = VecRestoreArray(v,&array);CHKERRQ(ierr); 209 ierr = VecDestroy(&v);CHKERRQ(ierr); 210 } else { 211 PetscMPIInt rank; 212 PetscInt nvGlobal, *offsets, myFirst, myLast; 213 214 ierr = PetscMalloc1(size+1,&offsets);CHKERRQ(ierr); 215 offsets[0] = 0; 216 ierr = MPI_Allgather(&numVertices,1,MPIU_INT,&offsets[1],1,MPIU_INT,comm);CHKERRMPI(ierr); 217 for (np = 2; np <= size; np++) { 218 offsets[np] += offsets[np-1]; 219 } 220 nvGlobal = offsets[size]; 221 ierr = MPI_Comm_rank(comm,&rank);CHKERRMPI(ierr); 222 myFirst = offsets[rank]; 223 myLast = offsets[rank + 1] - 1; 224 ierr = PetscFree(offsets);CHKERRQ(ierr); 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; 231 232 if (rem) { 233 firstLargePart = myFirst / pBig; 234 lastLargePart = myLast / pBig; 235 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 } 250 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); 254 255 PartStart = PetscMax(PartStart,myFirst); 256 PartEnd = PetscMin(PartEnd,myLast+1); 257 ierr = PetscSectionSetDof(partSection,np,PartEnd-PartStart);CHKERRQ(ierr); 258 } 259 } 260 } 261 } 262 ierr = PetscFree(tpwgts);CHKERRQ(ierr); 263 PetscFunctionReturn(0); 264 } 265 266 static PetscErrorCode PetscPartitionerInitialize_Simple(PetscPartitioner part) 267 { 268 PetscFunctionBegin; 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 PetscFunctionReturn(0); 275 } 276 277 /*MC 278 PETSCPARTITIONERSIMPLE = "simple" - A PetscPartitioner object 279 280 Level: intermediate 281 282 .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType() 283 M*/ 284 285 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Simple(PetscPartitioner part) 286 { 287 PetscPartitioner_Simple *p; 288 PetscErrorCode ierr; 289 290 PetscFunctionBegin; 291 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 292 ierr = PetscNewLog(part, &p);CHKERRQ(ierr); 293 p->gridDim = -1; 294 part->data = p; 295 296 ierr = PetscPartitionerInitialize_Simple(part);CHKERRQ(ierr); 297 PetscFunctionReturn(0); 298 } 299