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