1 #include <petsc/private/partitionerimpl.h> /*I "petscpartitioner.h" I*/ 2 3 typedef struct { 4 PetscSection section; /* Sizes for each partition */ 5 IS partition; /* Points in each partition */ 6 PetscBool random; /* Flag for a random partition */ 7 } PetscPartitioner_Shell; 8 9 static PetscErrorCode PetscPartitionerReset_Shell(PetscPartitioner part) 10 { 11 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 12 13 PetscFunctionBegin; 14 PetscCall(PetscSectionDestroy(&p->section)); 15 PetscCall(ISDestroy(&p->partition)); 16 PetscFunctionReturn(PETSC_SUCCESS); 17 } 18 19 static PetscErrorCode PetscPartitionerDestroy_Shell(PetscPartitioner part) 20 { 21 PetscFunctionBegin; 22 PetscCall(PetscPartitionerReset_Shell(part)); 23 PetscCall(PetscFree(part->data)); 24 PetscFunctionReturn(PETSC_SUCCESS); 25 } 26 27 static PetscErrorCode PetscPartitionerView_Shell_ASCII(PetscPartitioner part, PetscViewer viewer) 28 { 29 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 30 31 PetscFunctionBegin; 32 if (p->random) { 33 PetscCall(PetscViewerASCIIPushTab(viewer)); 34 PetscCall(PetscViewerASCIIPrintf(viewer, "using random partition\n")); 35 PetscCall(PetscViewerASCIIPopTab(viewer)); 36 } 37 PetscFunctionReturn(PETSC_SUCCESS); 38 } 39 40 static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer) 41 { 42 PetscBool iascii; 43 44 PetscFunctionBegin; 45 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 46 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 47 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii)); 48 if (iascii) PetscCall(PetscPartitionerView_Shell_ASCII(part, viewer)); 49 PetscFunctionReturn(PETSC_SUCCESS); 50 } 51 52 static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscPartitioner part, PetscOptionItems PetscOptionsObject) 53 { 54 PetscInt sizes[16], points[1024]; 55 PetscInt Npart = 16, Npoints = 1024; 56 PetscBool random = PETSC_FALSE, set, flgSizes, flgPoints; 57 PetscMPIInt rank; 58 59 PetscFunctionBegin; 60 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank)); 61 PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Shell Options"); 62 PetscCall(PetscPartitionerShellGetRandom(part, &random)); 63 PetscCall(PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set)); 64 if (set) PetscCall(PetscPartitionerShellSetRandom(part, random)); 65 PetscCall(PetscOptionsIntArray("-petscpartitioner_shell_sizes", "The size of each partition on rank 0", "PetscPartitionerShellSetPartition", sizes, &Npart, &flgSizes)); 66 PetscCall(PetscOptionsIntArray("-petscpartitioner_shell_points", "The points in each partition on rank 0", "PetscPartitionerShellSetPartition", points, &Npoints, &flgPoints)); 67 PetscCheck(!(flgSizes ^ flgPoints), PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Must specify both the partition sizes and points"); 68 if (flgSizes) { 69 PetscInt Np = 0; 70 71 for (PetscInt i = 0; i < Npart; ++i) Np += sizes[i]; 72 PetscCheck(Np == Npoints, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Number of input points %" PetscInt_FMT " != %" PetscInt_FMT " sum of partition sizes", Npoints, Np); 73 if (!rank) PetscCall(PetscPartitionerShellSetPartition(part, Npart, sizes, points)); 74 else { 75 PetscCall(PetscArrayzero(sizes, Npart)); 76 PetscCall(PetscPartitionerShellSetPartition(part, Npart, sizes, points)); 77 } 78 } 79 PetscOptionsHeadEnd(); 80 PetscFunctionReturn(PETSC_SUCCESS); 81 } 82 83 static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *partition) 84 { 85 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 86 PetscInt np; 87 88 PetscFunctionBegin; 89 if (p->random) { 90 PetscRandom r; 91 PetscInt *sizes, *points, v, p; 92 PetscMPIInt rank; 93 94 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank)); 95 PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &r)); 96 PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)nparts)); 97 PetscCall(PetscRandomSetFromOptions(r)); 98 PetscCall(PetscCalloc2(nparts, &sizes, numVertices, &points)); 99 for (v = 0; v < numVertices; ++v) points[v] = v; 100 for (p = 0; p < nparts; ++p) sizes[p] = numVertices / nparts + (PetscInt)(p < numVertices % nparts); 101 for (v = numVertices - 1; v > 0; --v) { 102 PetscReal val; 103 PetscInt w, tmp; 104 105 PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)(v + 1))); 106 PetscCall(PetscRandomGetValueReal(r, &val)); 107 w = PetscFloorReal(val); 108 tmp = points[v]; 109 points[v] = points[w]; 110 points[w] = tmp; 111 } 112 PetscCall(PetscRandomDestroy(&r)); 113 PetscCall(PetscPartitionerShellSetPartition(part, nparts, sizes, points)); 114 PetscCall(PetscFree2(sizes, points)); 115 } 116 PetscCheck(p->section, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Shell partitioner information not provided. Please call PetscPartitionerShellSetPartition()"); 117 PetscCall(PetscSectionGetChart(p->section, NULL, &np)); 118 PetscCheck(nparts == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of requested partitions %" PetscInt_FMT " != configured partitions %" PetscInt_FMT, nparts, np); 119 PetscCall(ISGetLocalSize(p->partition, &np)); 120 PetscCheck(numVertices == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of input vertices %" PetscInt_FMT " != configured vertices %" PetscInt_FMT, numVertices, np); 121 PetscCall(PetscSectionCopy(p->section, partSection)); 122 *partition = p->partition; 123 PetscCall(PetscObjectReference((PetscObject)p->partition)); 124 PetscFunctionReturn(PETSC_SUCCESS); 125 } 126 127 static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part) 128 { 129 PetscFunctionBegin; 130 part->noGraph = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */ 131 part->ops->view = PetscPartitionerView_Shell; 132 part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell; 133 part->ops->reset = PetscPartitionerReset_Shell; 134 part->ops->destroy = PetscPartitionerDestroy_Shell; 135 part->ops->partition = PetscPartitionerPartition_Shell; 136 PetscFunctionReturn(PETSC_SUCCESS); 137 } 138 139 /*MC 140 PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object 141 142 Level: intermediate 143 144 Options Database Keys: 145 . -petscpartitioner_shell_random - Use a random partition 146 147 .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()` 148 M*/ 149 150 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part) 151 { 152 PetscPartitioner_Shell *p; 153 154 PetscFunctionBegin; 155 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 156 PetscCall(PetscNew(&p)); 157 part->data = p; 158 159 PetscCall(PetscPartitionerInitialize_Shell(part)); 160 p->random = PETSC_FALSE; 161 PetscFunctionReturn(PETSC_SUCCESS); 162 } 163 164 /*@ 165 PetscPartitionerShellSetPartition - Set an artificial partition for a mesh 166 167 Collective 168 169 Input Parameters: 170 + part - The `PetscPartitioner` 171 . size - The number of partitions 172 . sizes - array of length size (or `NULL`) providing the number of points in each partition 173 - points - array of length sum(sizes) (may be `NULL` iff sizes is `NULL`), a permutation of the points that groups those assigned to each partition in order (i.e., partition 0 first, partition 1 next, etc.) 174 175 Level: developer 176 177 Note: 178 It is safe to free the sizes and points arrays after use in this routine. 179 180 .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()` 181 @*/ 182 PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[]) 183 { 184 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 185 PetscInt proc, numPoints; 186 187 PetscFunctionBegin; 188 PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL); 189 if (sizes) PetscAssertPointer(sizes, 3); 190 if (points) PetscAssertPointer(points, 4); 191 PetscCall(PetscSectionDestroy(&p->section)); 192 PetscCall(ISDestroy(&p->partition)); 193 PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)part), &p->section)); 194 PetscCall(PetscSectionSetChart(p->section, 0, size)); 195 if (sizes) { 196 for (proc = 0; proc < size; ++proc) PetscCall(PetscSectionSetDof(p->section, proc, sizes[proc])); 197 } 198 PetscCall(PetscSectionSetUp(p->section)); 199 PetscCall(PetscSectionGetStorageSize(p->section, &numPoints)); 200 PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)part), numPoints, points, PETSC_COPY_VALUES, &p->partition)); 201 PetscFunctionReturn(PETSC_SUCCESS); 202 } 203 204 /*@ 205 PetscPartitionerShellSetRandom - Set the flag to use a random partition 206 207 Collective 208 209 Input Parameters: 210 + part - The `PetscPartitioner` 211 - random - The flag to use a random partition 212 213 Level: intermediate 214 215 .seealso: `PetscPartitionerShellGetRandom()`, `PetscPartitionerCreate()` 216 @*/ 217 PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random) 218 { 219 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 220 221 PetscFunctionBegin; 222 PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL); 223 p->random = random; 224 PetscFunctionReturn(PETSC_SUCCESS); 225 } 226 227 /*@ 228 PetscPartitionerShellGetRandom - get the flag to use a random partition 229 230 Collective 231 232 Input Parameter: 233 . part - The `PetscPartitioner` 234 235 Output Parameter: 236 . random - The flag to use a random partition 237 238 Level: intermediate 239 240 .seealso: `PetscPartitionerShellSetRandom()`, `PetscPartitionerCreate()` 241 @*/ 242 PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random) 243 { 244 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 245 246 PetscFunctionBegin; 247 PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL); 248 PetscAssertPointer(random, 2); 249 *random = p->random; 250 PetscFunctionReturn(PETSC_SUCCESS); 251 } 252