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 PetscBool random = PETSC_FALSE, set; 55 56 PetscFunctionBegin; 57 PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Shell Options"); 58 PetscCall(PetscPartitionerShellGetRandom(part, &random)); 59 PetscCall(PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set)); 60 if (set) PetscCall(PetscPartitionerShellSetRandom(part, random)); 61 PetscOptionsHeadEnd(); 62 PetscFunctionReturn(PETSC_SUCCESS); 63 } 64 65 static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition) 66 { 67 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 68 PetscInt np; 69 70 PetscFunctionBegin; 71 if (p->random) { 72 PetscRandom r; 73 PetscInt *sizes, *points, v, p; 74 PetscMPIInt rank; 75 76 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank)); 77 PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &r)); 78 PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)nparts)); 79 PetscCall(PetscRandomSetFromOptions(r)); 80 PetscCall(PetscCalloc2(nparts, &sizes, numVertices, &points)); 81 for (v = 0; v < numVertices; ++v) points[v] = v; 82 for (p = 0; p < nparts; ++p) sizes[p] = numVertices / nparts + (PetscInt)(p < numVertices % nparts); 83 for (v = numVertices - 1; v > 0; --v) { 84 PetscReal val; 85 PetscInt w, tmp; 86 87 PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)(v + 1))); 88 PetscCall(PetscRandomGetValueReal(r, &val)); 89 w = PetscFloorReal(val); 90 tmp = points[v]; 91 points[v] = points[w]; 92 points[w] = tmp; 93 } 94 PetscCall(PetscRandomDestroy(&r)); 95 PetscCall(PetscPartitionerShellSetPartition(part, nparts, sizes, points)); 96 PetscCall(PetscFree2(sizes, points)); 97 } 98 PetscCheck(p->section, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Shell partitioner information not provided. Please call PetscPartitionerShellSetPartition()"); 99 PetscCall(PetscSectionGetChart(p->section, NULL, &np)); 100 PetscCheck(nparts == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of requested partitions %" PetscInt_FMT " != configured partitions %" PetscInt_FMT, nparts, np); 101 PetscCall(ISGetLocalSize(p->partition, &np)); 102 PetscCheck(numVertices == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of input vertices %" PetscInt_FMT " != configured vertices %" PetscInt_FMT, numVertices, np); 103 PetscCall(PetscSectionCopy(p->section, partSection)); 104 *partition = p->partition; 105 PetscCall(PetscObjectReference((PetscObject)p->partition)); 106 PetscFunctionReturn(PETSC_SUCCESS); 107 } 108 109 static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part) 110 { 111 PetscFunctionBegin; 112 part->noGraph = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */ 113 part->ops->view = PetscPartitionerView_Shell; 114 part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell; 115 part->ops->reset = PetscPartitionerReset_Shell; 116 part->ops->destroy = PetscPartitionerDestroy_Shell; 117 part->ops->partition = PetscPartitionerPartition_Shell; 118 PetscFunctionReturn(PETSC_SUCCESS); 119 } 120 121 /*MC 122 PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object 123 124 Level: intermediate 125 126 Options Database Keys: 127 . -petscpartitioner_shell_random - Use a random partition 128 129 .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()` 130 M*/ 131 132 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part) 133 { 134 PetscPartitioner_Shell *p; 135 136 PetscFunctionBegin; 137 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 138 PetscCall(PetscNew(&p)); 139 part->data = p; 140 141 PetscCall(PetscPartitionerInitialize_Shell(part)); 142 p->random = PETSC_FALSE; 143 PetscFunctionReturn(PETSC_SUCCESS); 144 } 145 146 /*@C 147 PetscPartitionerShellSetPartition - Set an artificial partition for a mesh 148 149 Collective 150 151 Input Parameters: 152 + part - The `PetscPartitioner` 153 . size - The number of partitions 154 . sizes - array of length size (or `NULL`) providing the number of points in each partition 155 - 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.) 156 157 Level: developer 158 159 Note: 160 It is safe to free the sizes and points arrays after use in this routine. 161 162 .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()` 163 @*/ 164 PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[]) 165 { 166 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 167 PetscInt proc, numPoints; 168 169 PetscFunctionBegin; 170 PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL); 171 if (sizes) PetscAssertPointer(sizes, 3); 172 if (points) PetscAssertPointer(points, 4); 173 PetscCall(PetscSectionDestroy(&p->section)); 174 PetscCall(ISDestroy(&p->partition)); 175 PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)part), &p->section)); 176 PetscCall(PetscSectionSetChart(p->section, 0, size)); 177 if (sizes) { 178 for (proc = 0; proc < size; ++proc) PetscCall(PetscSectionSetDof(p->section, proc, sizes[proc])); 179 } 180 PetscCall(PetscSectionSetUp(p->section)); 181 PetscCall(PetscSectionGetStorageSize(p->section, &numPoints)); 182 PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)part), numPoints, points, PETSC_COPY_VALUES, &p->partition)); 183 PetscFunctionReturn(PETSC_SUCCESS); 184 } 185 186 /*@ 187 PetscPartitionerShellSetRandom - Set the flag to use a random partition 188 189 Collective 190 191 Input Parameters: 192 + part - The `PetscPartitioner` 193 - random - The flag to use a random partition 194 195 Level: intermediate 196 197 .seealso: `PetscPartitionerShellGetRandom()`, `PetscPartitionerCreate()` 198 @*/ 199 PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random) 200 { 201 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 202 203 PetscFunctionBegin; 204 PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL); 205 p->random = random; 206 PetscFunctionReturn(PETSC_SUCCESS); 207 } 208 209 /*@ 210 PetscPartitionerShellGetRandom - get the flag to use a random partition 211 212 Collective 213 214 Input Parameter: 215 . part - The `PetscPartitioner` 216 217 Output Parameter: 218 . random - The flag to use a random partition 219 220 Level: intermediate 221 222 .seealso: `PetscPartitionerShellSetRandom()`, `PetscPartitionerCreate()` 223 @*/ 224 PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random) 225 { 226 PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data; 227 228 PetscFunctionBegin; 229 PetscValidHeaderSpecificType(part, PETSCPARTITIONER_CLASSID, 1, PETSCPARTITIONERSHELL); 230 PetscAssertPointer(random, 2); 231 *random = p->random; 232 PetscFunctionReturn(PETSC_SUCCESS); 233 } 234