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