1 #include <petscmat.h> 2 #include <petsc/private/partitionerimpl.h> /*I "petscpartitioner.h" I*/ 3 4 typedef struct { 5 MatPartitioning mp; 6 } PetscPartitioner_MatPartitioning; 7 8 static PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning(PetscPartitioner part, MatPartitioning *mp) 9 { 10 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data; 11 12 PetscFunctionBegin; 13 *mp = p->mp; 14 PetscFunctionReturn(PETSC_SUCCESS); 15 } 16 17 /*@ 18 PetscPartitionerMatPartitioningGetMatPartitioning - Get a `MatPartitioning` instance wrapped by this `PetscPartitioner`. 19 20 Not Collective 21 22 Input Parameter: 23 . part - The `PetscPartitioner` 24 25 Output Parameter: 26 . mp - The `MatPartitioning` 27 28 Level: developer 29 30 .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()` 31 @*/ 32 PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp) 33 { 34 PetscFunctionBegin; 35 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 36 PetscAssertPointer(mp, 2); 37 PetscUseMethod(part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", (PetscPartitioner, MatPartitioning *), (part, mp)); 38 PetscFunctionReturn(PETSC_SUCCESS); 39 } 40 41 static PetscErrorCode PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part) 42 { 43 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data; 44 45 PetscFunctionBegin; 46 PetscCall(MatPartitioningDestroy(&p->mp)); 47 PetscCall(PetscObjectComposeFunction((PetscObject)part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", NULL)); 48 PetscCall(PetscFree(part->data)); 49 PetscFunctionReturn(PETSC_SUCCESS); 50 } 51 52 static PetscErrorCode PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part, PetscViewer viewer) 53 { 54 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data; 55 PetscViewerFormat format; 56 57 PetscFunctionBegin; 58 PetscCall(PetscViewerGetFormat(viewer, &format)); 59 PetscCall(PetscViewerASCIIPrintf(viewer, "MatPartitioning Graph Partitioner:\n")); 60 PetscCall(PetscViewerASCIIPushTab(viewer)); 61 if (p->mp) PetscCall(MatPartitioningView(p->mp, viewer)); 62 PetscCall(PetscViewerASCIIPopTab(viewer)); 63 PetscFunctionReturn(PETSC_SUCCESS); 64 } 65 66 static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer) 67 { 68 PetscBool isascii; 69 70 PetscFunctionBegin; 71 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 72 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 73 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii)); 74 if (isascii) PetscCall(PetscPartitionerView_MatPartitioning_ASCII(part, viewer)); 75 PetscFunctionReturn(PETSC_SUCCESS); 76 } 77 78 static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscPartitioner part, PetscOptionItems PetscOptionsObject) 79 { 80 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data; 81 82 PetscFunctionBegin; 83 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)p->mp, ((PetscObject)part)->prefix)); 84 PetscCall(MatPartitioningSetFromOptions(p->mp)); 85 PetscFunctionReturn(PETSC_SUCCESS); 86 } 87 88 static PetscErrorCode PetscPartitionerPartition_MatPartitioning(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *is) 89 { 90 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data; 91 Mat matadj; 92 IS is1, is2, is3; 93 PetscReal *tpwgts = NULL; 94 PetscInt numVerticesGlobal, numEdges; 95 PetscInt *i, *j, *vwgt = NULL; 96 MPI_Comm comm; 97 98 PetscFunctionBegin; 99 PetscCall(PetscObjectGetComm((PetscObject)part, &comm)); 100 101 /* TODO: MatCreateMPIAdj should maybe take global number of ROWS */ 102 /* TODO: And vertex distribution in PetscPartitionerPartition_ParMetis should be done using PetscSplitOwnership */ 103 numVerticesGlobal = PETSC_DECIDE; 104 PetscCall(PetscSplitOwnership(comm, &numVertices, &numVerticesGlobal)); 105 106 /* copy arrays to avoid memory errors because MatMPIAdjSetPreallocation copies just pointers */ 107 numEdges = start[numVertices]; 108 PetscCall(PetscMalloc1(numVertices + 1, &i)); 109 PetscCall(PetscMalloc1(numEdges, &j)); 110 PetscCall(PetscArraycpy(i, start, numVertices + 1)); 111 PetscCall(PetscArraycpy(j, adjacency, numEdges)); 112 113 /* construct the adjacency matrix */ 114 PetscCall(MatCreateMPIAdj(comm, numVertices, numVerticesGlobal, i, j, NULL, &matadj)); 115 PetscCall(MatPartitioningSetAdjacency(p->mp, matadj)); 116 PetscCall(MatPartitioningSetNParts(p->mp, nparts)); 117 118 /* calculate partition weights */ 119 if (targetSection) { 120 PetscReal sumt; 121 PetscInt p; 122 123 sumt = 0.0; 124 PetscCall(PetscMalloc1(nparts, &tpwgts)); 125 for (p = 0; p < nparts; ++p) { 126 PetscInt tpd; 127 128 PetscCall(PetscSectionGetDof(targetSection, p, &tpd)); 129 sumt += tpd; 130 tpwgts[p] = tpd; 131 } 132 if (sumt) { /* METIS/ParMETIS do not like exactly zero weight */ 133 for (p = 0, sumt = 0.0; p < nparts; ++p) { 134 tpwgts[p] = PetscMax(tpwgts[p], PETSC_SMALL); 135 sumt += tpwgts[p]; 136 } 137 for (p = 0; p < nparts; ++p) tpwgts[p] /= sumt; 138 for (p = 0, sumt = 0.0; p < nparts - 1; ++p) sumt += tpwgts[p]; 139 tpwgts[nparts - 1] = 1. - sumt; 140 } else { 141 PetscCall(PetscFree(tpwgts)); 142 } 143 } 144 PetscCall(MatPartitioningSetPartitionWeights(p->mp, tpwgts)); 145 146 /* calculate vertex weights */ 147 if (vertSection) { 148 PetscInt v; 149 150 PetscCall(PetscMalloc1(numVertices, &vwgt)); 151 for (v = 0; v < numVertices; ++v) PetscCall(PetscSectionGetDof(vertSection, v, &vwgt[v])); 152 } 153 PetscCall(MatPartitioningSetVertexWeights(p->mp, vwgt)); 154 155 /* apply the partitioning */ 156 PetscCall(MatPartitioningApply(p->mp, &is1)); 157 158 /* construct the PetscSection */ 159 { 160 PetscInt v; 161 const PetscInt *assignment_arr; 162 163 PetscCall(ISGetIndices(is1, &assignment_arr)); 164 for (v = 0; v < numVertices; ++v) PetscCall(PetscSectionAddDof(partSection, assignment_arr[v], 1)); 165 PetscCall(ISRestoreIndices(is1, &assignment_arr)); 166 } 167 168 /* convert assignment IS to global numbering IS */ 169 PetscCall(ISPartitioningToNumbering(is1, &is2)); 170 PetscCall(ISDestroy(&is1)); 171 172 /* renumber IS into local numbering */ 173 PetscCall(ISOnComm(is2, PETSC_COMM_SELF, PETSC_USE_POINTER, &is1)); 174 PetscCall(ISRenumber(is1, NULL, NULL, &is3)); 175 PetscCall(ISDestroy(&is1)); 176 PetscCall(ISDestroy(&is2)); 177 178 /* invert IS */ 179 PetscCall(ISSetPermutation(is3)); 180 PetscCall(ISInvertPermutation(is3, numVertices, &is1)); 181 PetscCall(ISDestroy(&is3)); 182 183 PetscCall(MatDestroy(&matadj)); 184 *is = is1; 185 PetscFunctionReturn(PETSC_SUCCESS); 186 } 187 188 static PetscErrorCode PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part) 189 { 190 PetscFunctionBegin; 191 part->ops->view = PetscPartitionerView_MatPartitioning; 192 part->ops->setfromoptions = PetscPartitionerSetFromOptions_MatPartitioning; 193 part->ops->destroy = PetscPartitionerDestroy_MatPartitioning; 194 part->ops->partition = PetscPartitionerPartition_MatPartitioning; 195 PetscCall(PetscObjectComposeFunction((PetscObject)part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning)); 196 PetscFunctionReturn(PETSC_SUCCESS); 197 } 198 199 /*MC 200 PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object 201 202 Level: developer 203 204 .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()` 205 M*/ 206 207 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part) 208 { 209 PetscPartitioner_MatPartitioning *p; 210 211 PetscFunctionBegin; 212 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 213 PetscCall(PetscNew(&p)); 214 part->data = p; 215 PetscCall(PetscPartitionerInitialize_MatPartitioning(part)); 216 PetscCall(MatPartitioningCreate(PetscObjectComm((PetscObject)part), &p->mp)); 217 PetscFunctionReturn(PETSC_SUCCESS); 218 } 219