Date: 2022-03-04T17:38:49+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;

     13:   *mp = p->mp;
     14:   return(0);
     15: }

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

     20:   Not Collective

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

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

     28:   Level: developer

     30: .seealso DMPlexDistribute(), PetscPartitionerCreate()
     31: @*/
     32: PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp)
     33: {
     34:   PetscErrorCode          ierr;

     39:   PetscUseMethod(part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",(PetscPartitioner,MatPartitioning*),(part,mp));
     40:   return(0);
     41: }

     43: static PetscErrorCode PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part)
     44: {
     45:   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
     46:   PetscErrorCode                    ierr;

     49:   MatPartitioningDestroy(&p->mp);
     50:   PetscFree(part->data);
     51:   return(0);
     52: }

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

     61:   PetscViewerGetFormat(viewer, &format);
     62:   PetscViewerASCIIPrintf(viewer, "MatPartitioning Graph Partitioner:\n");
     63:   PetscViewerASCIIPushTab(viewer);
     64:   if (p->mp) {MatPartitioningView(p->mp, viewer);}
     65:   PetscViewerASCIIPopTab(viewer);
     66:   return(0);
     67: }

     69: static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer)
     70: {
     71:   PetscBool      iascii;

     77:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
     78:   if (iascii) {PetscPartitionerView_MatPartitioning_ASCII(part, viewer);}
     79:   return(0);
     80: }

     82: static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscOptionItems *PetscOptionsObject, PetscPartitioner part)
     83: {
     84:   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
     85:   PetscErrorCode                    ierr;

     88:   PetscObjectSetOptionsPrefix((PetscObject)p->mp,((PetscObject)part)->prefix);
     89:   MatPartitioningSetFromOptions(p->mp);
     90:   return(0);
     91: }

     93: static PetscErrorCode PetscPartitionerPartition_MatPartitioning(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *is)
     94: {
     95:   PetscPartitioner_MatPartitioning  *p = (PetscPartitioner_MatPartitioning *) part->data;
     96:   Mat                               matadj;
     97:   IS                                is1, is2, is3;
     98:   PetscReal                         *tpwgts = NULL;
     99:   PetscInt                          numVerticesGlobal, numEdges;
    100:   PetscInt                          *i, *j, *vwgt = NULL;
    101:   MPI_Comm                          comm;
    102:   PetscErrorCode                    ierr;

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

    107:   /* TODO: MatCreateMPIAdj should maybe take global number of ROWS */
    108:   /* TODO: And vertex distribution in PetscPartitionerPartition_ParMetis should be done using PetscSplitOwnership */
    109:   numVerticesGlobal = PETSC_DECIDE;
    110:   PetscSplitOwnership(comm, &numVertices, &numVerticesGlobal);

    112:   /* copy arrays to avoid memory errors because MatMPIAdjSetPreallocation copies just pointers */
    113:   numEdges = start[numVertices];
    114:   PetscMalloc1(numVertices+1, &i);
    115:   PetscMalloc1(numEdges, &j);
    116:   PetscArraycpy(i, start, numVertices+1);
    117:   PetscArraycpy(j, adjacency, numEdges);

    119:   /* construct the adjacency matrix */
    120:   MatCreateMPIAdj(comm, numVertices, numVerticesGlobal, i, j, NULL, &matadj);
    121:   MatPartitioningSetAdjacency(p->mp, matadj);
    122:   MatPartitioningSetNParts(p->mp, nparts);

    124:   /* calculate partition weights */
    125:   if (targetSection) {
    126:     PetscReal sumt;
    127:     PetscInt  p;

    129:     sumt = 0.0;
    130:     PetscMalloc1(nparts,&tpwgts);
    131:     for (p = 0; p < nparts; ++p) {
    132:       PetscInt tpd;

    134:       PetscSectionGetDof(targetSection,p,&tpd);
    135:       sumt += tpd;
    136:       tpwgts[p] = tpd;
    137:     }
    138:     if (sumt) { /* METIS/ParMETIS do not like exactly zero weight */
    139:       for (p = 0, sumt = 0.0; p < nparts; ++p) {
    140:         tpwgts[p] = PetscMax(tpwgts[p],PETSC_SMALL);
    141:         sumt += tpwgts[p];
    142:       }
    143:       for (p = 0; p < nparts; ++p) tpwgts[p] /= sumt;
    144:       for (p = 0, sumt = 0.0; p < nparts-1; ++p) sumt += tpwgts[p];
    145:       tpwgts[nparts - 1] = 1. - sumt;
    146:     } else {
    147:       PetscFree(tpwgts);
    148:     }
    149:   }
    150:   MatPartitioningSetPartitionWeights(p->mp, tpwgts);

    152:   /* calculate vertex weights */
    153:   if (vertSection) {
    154:     PetscInt v;

    156:     PetscMalloc1(numVertices,&vwgt);
    157:     for (v = 0; v < numVertices; ++v) {
    158:       PetscSectionGetDof(vertSection, v, &vwgt[v]);
    159:     }
    160:   }
    161:   MatPartitioningSetVertexWeights(p->mp, vwgt);

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

    166:   /* construct the PetscSection */
    167:   {
    168:     PetscInt v;
    169:     const PetscInt *assignment_arr;

    171:     ISGetIndices(is1, &assignment_arr);
    172:     for (v = 0; v < numVertices; ++v) {PetscSectionAddDof(partSection, assignment_arr[v], 1);}
    173:     ISRestoreIndices(is1, &assignment_arr);
    174:   }

    176:   /* convert assignment IS to global numbering IS */
    177:   ISPartitioningToNumbering(is1, &is2);
    178:   ISDestroy(&is1);

    180:   /* renumber IS into local numbering */
    181:   ISOnComm(is2, PETSC_COMM_SELF, PETSC_USE_POINTER, &is1);
    182:   ISRenumber(is1, NULL, NULL, &is3);
    183:   ISDestroy(&is1);
    184:   ISDestroy(&is2);

    186:   /* invert IS */
    187:   ISSetPermutation(is3);
    188:   ISInvertPermutation(is3, numVertices, &is1);
    189:   ISDestroy(&is3);

    191:   MatDestroy(&matadj);
    192:   *is = is1;
    193:   return(0);
    194: }

    196: static PetscErrorCode PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part)
    197: {
    198:   PetscErrorCode                    ierr;

    201:   part->ops->view           = PetscPartitionerView_MatPartitioning;
    202:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_MatPartitioning;
    203:   part->ops->destroy        = PetscPartitionerDestroy_MatPartitioning;
    204:   part->ops->partition      = PetscPartitionerPartition_MatPartitioning;
    205:   PetscObjectComposeFunction((PetscObject)part,"PetscPartitionerMatPartitioningGetMatPartitioning_C",PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning);
    206:   return(0);
    207: }

    209: /*MC
    210:   PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object

    212:   Level: developer

    214: .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
    215: M*/

    217: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)
    218: {
    219:   PetscPartitioner_MatPartitioning  *p;
    220:   PetscErrorCode                    ierr;

    224:   PetscNewLog(part, &p);
    225:   part->data = p;
    226:   PetscPartitionerInitialize_MatPartitioning(part);
    227:   MatPartitioningCreate(PetscObjectComm((PetscObject)part), &p->mp);
    228:   return(0);
    229: }