Date: 2022-03-31T04:47:55+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:   *mp = p->mp;
     13:   return 0;
     14: }

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

     19:   Not Collective

     21:   Input Parameters:
     22: . part     - The PetscPartitioner

     24:   Output Parameters:
     25: . mp       - The MatPartitioning

     27:   Level: developer

     29: .seealso DMPlexDistribute(), PetscPartitionerCreate()
     30: @*/
     31: PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp)
     32: {
     35:   PetscUseMethod(part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",(PetscPartitioner,MatPartitioning*),(part,mp));
     36:   return 0;
     37: }

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

     43:   MatPartitioningDestroy(&p->mp);
     44:   PetscFree(part->data);
     45:   return 0;
     46: }

     48: static PetscErrorCode PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part, PetscViewer viewer)
     49: {
     50:   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
     51:   PetscViewerFormat                 format;

     53:   PetscViewerGetFormat(viewer, &format);
     54:   PetscViewerASCIIPrintf(viewer, "MatPartitioning Graph Partitioner:\n");
     55:   PetscViewerASCIIPushTab(viewer);
     56:   if (p->mp) MatPartitioningView(p->mp, viewer);
     57:   PetscViewerASCIIPopTab(viewer);
     58:   return 0;
     59: }

     61: static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer)
     62: {
     63:   PetscBool      iascii;

     67:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
     68:   if (iascii) PetscPartitionerView_MatPartitioning_ASCII(part, viewer);
     69:   return 0;
     70: }

     72: static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscOptionItems *PetscOptionsObject, PetscPartitioner part)
     73: {
     74:   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;

     76:   PetscObjectSetOptionsPrefix((PetscObject)p->mp,((PetscObject)part)->prefix);
     77:   MatPartitioningSetFromOptions(p->mp);
     78:   return 0;
     79: }

     81: static PetscErrorCode PetscPartitionerPartition_MatPartitioning(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *is)
     82: {
     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;

     91:   PetscObjectGetComm((PetscObject)part, &comm);

     93:   /* TODO: MatCreateMPIAdj should maybe take global number of ROWS */
     94:   /* TODO: And vertex distribution in PetscPartitionerPartition_ParMetis should be done using PetscSplitOwnership */
     95:   numVerticesGlobal = PETSC_DECIDE;
     96:   PetscSplitOwnership(comm, &numVertices, &numVerticesGlobal);

     98:   /* copy arrays to avoid memory errors because MatMPIAdjSetPreallocation copies just pointers */
     99:   numEdges = start[numVertices];
    100:   PetscMalloc1(numVertices+1, &i);
    101:   PetscMalloc1(numEdges, &j);
    102:   PetscArraycpy(i, start, numVertices+1);
    103:   PetscArraycpy(j, adjacency, numEdges);

    105:   /* construct the adjacency matrix */
    106:   MatCreateMPIAdj(comm, numVertices, numVerticesGlobal, i, j, NULL, &matadj);
    107:   MatPartitioningSetAdjacency(p->mp, matadj);
    108:   MatPartitioningSetNParts(p->mp, nparts);

    110:   /* calculate partition weights */
    111:   if (targetSection) {
    112:     PetscReal sumt;
    113:     PetscInt  p;

    115:     sumt = 0.0;
    116:     PetscMalloc1(nparts,&tpwgts);
    117:     for (p = 0; p < nparts; ++p) {
    118:       PetscInt tpd;

    120:       PetscSectionGetDof(targetSection,p,&tpd);
    121:       sumt += tpd;
    122:       tpwgts[p] = tpd;
    123:     }
    124:     if (sumt) { /* METIS/ParMETIS do not like exactly zero weight */
    125:       for (p = 0, sumt = 0.0; p < nparts; ++p) {
    126:         tpwgts[p] = PetscMax(tpwgts[p],PETSC_SMALL);
    127:         sumt += tpwgts[p];
    128:       }
    129:       for (p = 0; p < nparts; ++p) tpwgts[p] /= sumt;
    130:       for (p = 0, sumt = 0.0; p < nparts-1; ++p) sumt += tpwgts[p];
    131:       tpwgts[nparts - 1] = 1. - sumt;
    132:     } else {
    133:       PetscFree(tpwgts);
    134:     }
    135:   }
    136:   MatPartitioningSetPartitionWeights(p->mp, tpwgts);

    138:   /* calculate vertex weights */
    139:   if (vertSection) {
    140:     PetscInt v;

    142:     PetscMalloc1(numVertices,&vwgt);
    143:     for (v = 0; v < numVertices; ++v) {
    144:       PetscSectionGetDof(vertSection, v, &vwgt[v]);
    145:     }
    146:   }
    147:   MatPartitioningSetVertexWeights(p->mp, vwgt);

    149:   /* apply the partitioning */
    150:   MatPartitioningApply(p->mp, &is1);

    152:   /* construct the PetscSection */
    153:   {
    154:     PetscInt v;
    155:     const PetscInt *assignment_arr;

    157:     ISGetIndices(is1, &assignment_arr);
    158:     for (v = 0; v < numVertices; ++v) PetscSectionAddDof(partSection, assignment_arr[v], 1);
    159:     ISRestoreIndices(is1, &assignment_arr);
    160:   }

    162:   /* convert assignment IS to global numbering IS */
    163:   ISPartitioningToNumbering(is1, &is2);
    164:   ISDestroy(&is1);

    166:   /* renumber IS into local numbering */
    167:   ISOnComm(is2, PETSC_COMM_SELF, PETSC_USE_POINTER, &is1);
    168:   ISRenumber(is1, NULL, NULL, &is3);
    169:   ISDestroy(&is1);
    170:   ISDestroy(&is2);

    172:   /* invert IS */
    173:   ISSetPermutation(is3);
    174:   ISInvertPermutation(is3, numVertices, &is1);
    175:   ISDestroy(&is3);

    177:   MatDestroy(&matadj);
    178:   *is = is1;
    179:   return 0;
    180: }

    182: static PetscErrorCode PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part)
    183: {
    184:   part->ops->view           = PetscPartitionerView_MatPartitioning;
    185:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_MatPartitioning;
    186:   part->ops->destroy        = PetscPartitionerDestroy_MatPartitioning;
    187:   part->ops->partition      = PetscPartitionerPartition_MatPartitioning;
    188:   PetscObjectComposeFunction((PetscObject)part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning);
    189:   return 0;
    190: }

    192: /*MC
    193:   PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object

    195:   Level: developer

    197: .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
    198: M*/

    200: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)
    201: {
    202:   PetscPartitioner_MatPartitioning  *p;

    205:   PetscNewLog(part, &p);
    206:   part->data = p;
    207:   PetscPartitionerInitialize_MatPartitioning(part);
    208:   MatPartitioningCreate(PetscObjectComm((PetscObject)part), &p->mp);
    209:   return 0;
    210: }