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(0); 15 } 16 17 /*@C 18 PetscPartitionerMatPartitioningGetMatPartitioning - Get a MatPartitioning instance wrapped by this PetscPartitioner. 19 20 Not Collective 21 22 Input Parameters: 23 . part - The PetscPartitioner 24 25 Output Parameters: 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 PetscValidPointer(mp, 2); 37 PetscUseMethod(part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",(PetscPartitioner,MatPartitioning*),(part,mp)); 38 PetscFunctionReturn(0); 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(0); 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(0); 64 } 65 66 static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer) 67 { 68 PetscBool iascii; 69 70 PetscFunctionBegin; 71 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 72 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 73 PetscCall(PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii)); 74 if (iascii) PetscCall(PetscPartitionerView_MatPartitioning_ASCII(part, viewer)); 75 PetscFunctionReturn(0); 76 } 77 78 static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscOptionItems *PetscOptionsObject, PetscPartitioner part) 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(0); 86 } 87 88 static PetscErrorCode PetscPartitionerPartition_MatPartitioning(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, 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) { 152 PetscCall(PetscSectionGetDof(vertSection, v, &vwgt[v])); 153 } 154 } 155 PetscCall(MatPartitioningSetVertexWeights(p->mp, vwgt)); 156 157 /* apply the partitioning */ 158 PetscCall(MatPartitioningApply(p->mp, &is1)); 159 160 /* construct the PetscSection */ 161 { 162 PetscInt v; 163 const PetscInt *assignment_arr; 164 165 PetscCall(ISGetIndices(is1, &assignment_arr)); 166 for (v = 0; v < numVertices; ++v) PetscCall(PetscSectionAddDof(partSection, assignment_arr[v], 1)); 167 PetscCall(ISRestoreIndices(is1, &assignment_arr)); 168 } 169 170 /* convert assignment IS to global numbering IS */ 171 PetscCall(ISPartitioningToNumbering(is1, &is2)); 172 PetscCall(ISDestroy(&is1)); 173 174 /* renumber IS into local numbering */ 175 PetscCall(ISOnComm(is2, PETSC_COMM_SELF, PETSC_USE_POINTER, &is1)); 176 PetscCall(ISRenumber(is1, NULL, NULL, &is3)); 177 PetscCall(ISDestroy(&is1)); 178 PetscCall(ISDestroy(&is2)); 179 180 /* invert IS */ 181 PetscCall(ISSetPermutation(is3)); 182 PetscCall(ISInvertPermutation(is3, numVertices, &is1)); 183 PetscCall(ISDestroy(&is3)); 184 185 PetscCall(MatDestroy(&matadj)); 186 *is = is1; 187 PetscFunctionReturn(0); 188 } 189 190 static PetscErrorCode PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part) 191 { 192 PetscFunctionBegin; 193 part->ops->view = PetscPartitionerView_MatPartitioning; 194 part->ops->setfromoptions = PetscPartitionerSetFromOptions_MatPartitioning; 195 part->ops->destroy = PetscPartitionerDestroy_MatPartitioning; 196 part->ops->partition = PetscPartitionerPartition_MatPartitioning; 197 PetscCall(PetscObjectComposeFunction((PetscObject)part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning)); 198 PetscFunctionReturn(0); 199 } 200 201 /*MC 202 PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object 203 204 Level: developer 205 206 .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()` 207 M*/ 208 209 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part) 210 { 211 PetscPartitioner_MatPartitioning *p; 212 213 PetscFunctionBegin; 214 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1); 215 PetscCall(PetscNewLog(part, &p)); 216 part->data = p; 217 PetscCall(PetscPartitionerInitialize_MatPartitioning(part)); 218 PetscCall(MatPartitioningCreate(PetscObjectComm((PetscObject)part), &p->mp)); 219 PetscFunctionReturn(0); 220 } 221