Date: 2026-06-01T23:13:23+00:00

Actual source code: partmatpart.c

      1: #include <petscmat.h>
      2: #include <petsc/private/partitionerimpl.h>

      4: typedef struct {
      5:   MatPartitioning mp;
      6: } PetscPartitioner_MatPartitioning;

      8: static PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning(PetscPartitioner part, MatPartitioning *mp)
      9: {
     10:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;

     12:   PetscFunctionBegin;
     13:   *mp = p->mp;
     14:   PetscFunctionReturn(PETSC_SUCCESS);
     15: }

     17: /*@
     18:   PetscPartitionerMatPartitioningGetMatPartitioning - Get a `MatPartitioning` instance wrapped by this `PetscPartitioner`.

     20:   Not Collective

     22:   Input Parameter:
     23: . part - The `PetscPartitioner`

     25:   Output Parameter:
     26: . mp - The `MatPartitioning`

     28:   Level: developer

     30: .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()`
     31: @*/
     32: PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp)
     33: {
     34:   PetscFunctionBegin;
     36:   PetscAssertPointer(mp, 2);
     37:   PetscUseMethod(part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", (PetscPartitioner, MatPartitioning *), (part, mp));
     38:   PetscFunctionReturn(PETSC_SUCCESS);
     39: }

     41: static PetscErrorCode PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part)
     42: {
     43:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;

     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: }

     52: static PetscErrorCode PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part, PetscViewer viewer)
     53: {
     54:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
     55:   PetscViewerFormat                 format;

     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: }

     66: static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer)
     67: {
     68:   PetscBool isascii;

     70:   PetscFunctionBegin;
     73:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
     74:   if (isascii) PetscCall(PetscPartitionerView_MatPartitioning_ASCII(part, viewer));
     75:   PetscFunctionReturn(PETSC_SUCCESS);
     76: }

     78: static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscPartitioner part, PetscOptionItems PetscOptionsObject)
     79: {
     80:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;

     82:   PetscFunctionBegin;
     83:   PetscCall(PetscObjectSetOptionsPrefix((PetscObject)p->mp, ((PetscObject)part)->prefix));
     84:   PetscCall(MatPartitioningSetFromOptions(p->mp));
     85:   PetscFunctionReturn(PETSC_SUCCESS);
     86: }

     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;

     98:   PetscFunctionBegin;
     99:   PetscCall(PetscObjectGetComm((PetscObject)part, &comm));

    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));

    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));

    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));

    118:   /* calculate partition weights */
    119:   if (targetSection) {
    120:     PetscReal sumt;
    121:     PetscInt  p;

    123:     sumt = 0.0;
    124:     PetscCall(PetscMalloc1(nparts, &tpwgts));
    125:     for (p = 0; p < nparts; ++p) {
    126:       PetscInt tpd;

    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));

    146:   /* calculate vertex weights */
    147:   if (vertSection) {
    148:     PetscInt v;

    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));

    155:   /* apply the partitioning */
    156:   PetscCall(MatPartitioningApply(p->mp, &is1));

    158:   /* construct the PetscSection */
    159:   {
    160:     PetscInt        v;
    161:     const PetscInt *assignment_arr;

    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:   }

    168:   /* convert assignment IS to global numbering IS */
    169:   PetscCall(ISPartitioningToNumbering(is1, &is2));
    170:   PetscCall(ISDestroy(&is1));

    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));

    178:   /* invert IS */
    179:   PetscCall(ISSetPermutation(is3));
    180:   PetscCall(ISInvertPermutation(is3, numVertices, &is1));
    181:   PetscCall(ISDestroy(&is3));

    183:   PetscCall(MatDestroy(&matadj));
    184:   *is = is1;
    185:   PetscFunctionReturn(PETSC_SUCCESS);
    186: }

    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: }

    199: /*MC
    200:   PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object

    202:   Level: developer

    204: .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
    205: M*/

    207: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)
    208: {
    209:   PetscPartitioner_MatPartitioning *p;

    211:   PetscFunctionBegin;
    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: }