Date: 2022-03-04T17:38:49+00:00

Actual source code: partshell.c

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

      3: typedef struct {
      4:   PetscSection section;   /* Sizes for each partition */
      5:   IS           partition; /* Points in each partition */
      6:   PetscBool    random;    /* Flag for a random partition */
      7: } PetscPartitioner_Shell;

      9: static PetscErrorCode PetscPartitionerReset_Shell(PetscPartitioner part)
     10: {
     11:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;
     12:   PetscErrorCode          ierr;

     15:   PetscSectionDestroy(&p->section);
     16:   ISDestroy(&p->partition);
     17:   return(0);
     18: }

     20: static PetscErrorCode PetscPartitionerDestroy_Shell(PetscPartitioner part)
     21: {

     25:   PetscPartitionerReset_Shell(part);
     26:   PetscFree(part->data);
     27:   return(0);
     28: }

     30: static PetscErrorCode PetscPartitionerView_Shell_ASCII(PetscPartitioner part, PetscViewer viewer)
     31: {
     32:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;
     33:   PetscErrorCode          ierr;

     36:   if (p->random) {
     37:     PetscViewerASCIIPushTab(viewer);
     38:     PetscViewerASCIIPrintf(viewer, "using random partition\n");
     39:     PetscViewerASCIIPopTab(viewer);
     40:   }
     41:   return(0);
     42: }

     44: static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer)
     45: {
     46:   PetscBool      iascii;

     52:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
     53:   if (iascii) {PetscPartitionerView_Shell_ASCII(part, viewer);}
     54:   return(0);
     55: }

     57: static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscOptionItems *PetscOptionsObject, PetscPartitioner part)
     58: {
     59:   PetscBool      random = PETSC_FALSE, set;

     63:   PetscOptionsHead(PetscOptionsObject, "PetscPartitioner Shell Options");
     64:   PetscPartitionerShellGetRandom(part, &random);
     65:   PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set);
     66:   if (set) {PetscPartitionerShellSetRandom(part, random);}
     67:   PetscOptionsTail();
     68:   return(0);
     69: }

     71: static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
     72: {
     73:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;
     74:   PetscInt                np;
     75:   PetscErrorCode          ierr;

     78:   if (p->random) {
     79:     PetscRandom r;
     80:     PetscInt   *sizes, *points, v, p;
     81:     PetscMPIInt rank;

     83:     MPI_Comm_rank(PetscObjectComm((PetscObject) part), &rank);
     84:     PetscRandomCreate(PETSC_COMM_SELF, &r);
     85:     PetscRandomSetInterval(r, 0.0, (PetscScalar) nparts);
     86:     PetscRandomSetFromOptions(r);
     87:     PetscCalloc2(nparts, &sizes, numVertices, &points);
     88:     for (v = 0; v < numVertices; ++v) {points[v] = v;}
     89:     for (p = 0; p < nparts; ++p) {sizes[p] = numVertices/nparts + (PetscInt) (p < numVertices % nparts);}
     90:     for (v = numVertices-1; v > 0; --v) {
     91:       PetscReal val;
     92:       PetscInt  w, tmp;

     94:       PetscRandomSetInterval(r, 0.0, (PetscScalar) (v+1));
     95:       PetscRandomGetValueReal(r, &val);
     96:       w    = PetscFloorReal(val);
     97:       tmp       = points[v];
     98:       points[v] = points[w];
     99:       points[w] = tmp;
    100:     }
    101:     PetscRandomDestroy(&r);
    102:     PetscPartitionerShellSetPartition(part, nparts, sizes, points);
    103:     PetscFree2(sizes, points);
    104:   }
    105:   if (!p->section) SETERRQ(PetscObjectComm((PetscObject) part), PETSC_ERR_ARG_WRONG, "Shell partitioner information not provided. Please call PetscPartitionerShellSetPartition()");
    106:   PetscSectionGetChart(p->section, NULL, &np);
    107:   if (nparts != np) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of requested partitions %d != configured partitions %d", nparts, np);
    108:   ISGetLocalSize(p->partition, &np);
    109:   if (numVertices != np) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of input vertices %d != configured vertices %d", numVertices, np);
    110:   PetscSectionCopy(p->section, partSection);
    111:   *partition = p->partition;
    112:   PetscObjectReference((PetscObject) p->partition);
    113:   return(0);
    114: }

    116: static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part)
    117: {
    119:   part->noGraph             = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */
    120:   part->ops->view           = PetscPartitionerView_Shell;
    121:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell;
    122:   part->ops->reset          = PetscPartitionerReset_Shell;
    123:   part->ops->destroy        = PetscPartitionerDestroy_Shell;
    124:   part->ops->partition      = PetscPartitionerPartition_Shell;
    125:   return(0);
    126: }

    128: /*MC
    129:   PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object

    131:   Level: intermediate

    133:   Options Database Keys:
    134: .  -petscpartitioner_shell_random - Use a random partition

    136: .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
    137: M*/

    139: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part)
    140: {
    141:   PetscPartitioner_Shell *p;
    142:   PetscErrorCode          ierr;

    146:   PetscNewLog(part, &p);
    147:   part->data = p;

    149:   PetscPartitionerInitialize_Shell(part);
    150:   p->random = PETSC_FALSE;
    151:   return(0);
    152: }

    154: /*@C
    155:   PetscPartitionerShellSetPartition - Set an artifical partition for a mesh

    157:   Collective on PetscPartitioner

    159:   Input Parameters:
    160: + part   - The PetscPartitioner
    161: . size   - The number of partitions
    162: . sizes  - array of length size (or NULL) providing the number of points in each partition
    163: - points - array of length sum(sizes) (may be NULL iff sizes is NULL), a permutation of the points that groups those assigned to each partition in order (i.e., partition 0 first, partition 1 next, etc.)

    165:   Level: developer

    167:   Notes:
    168:     It is safe to free the sizes and points arrays after use in this routine.

    170: .seealso DMPlexDistribute(), PetscPartitionerCreate()
    171: @*/
    172: PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[])
    173: {
    174:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;
    175:   PetscInt                proc, numPoints;
    176:   PetscErrorCode          ierr;

    182:   PetscSectionDestroy(&p->section);
    183:   ISDestroy(&p->partition);
    184:   PetscSectionCreate(PetscObjectComm((PetscObject) part), &p->section);
    185:   PetscSectionSetChart(p->section, 0, size);
    186:   if (sizes) {
    187:     for (proc = 0; proc < size; ++proc) {
    188:       PetscSectionSetDof(p->section, proc, sizes[proc]);
    189:     }
    190:   }
    191:   PetscSectionSetUp(p->section);
    192:   PetscSectionGetStorageSize(p->section, &numPoints);
    193:   ISCreateGeneral(PetscObjectComm((PetscObject) part), numPoints, points, PETSC_COPY_VALUES, &p->partition);
    194:   return(0);
    195: }

    197: /*@
    198:   PetscPartitionerShellSetRandom - Set the flag to use a random partition

    200:   Collective on PetscPartitioner

    202:   Input Parameters:
    203: + part   - The PetscPartitioner
    204: - random - The flag to use a random partition

    206:   Level: intermediate

    208: .seealso PetscPartitionerShellGetRandom(), PetscPartitionerCreate()
    209: @*/
    210: PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random)
    211: {
    212:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;

    216:   p->random = random;
    217:   return(0);
    218: }

    220: /*@
    221:   PetscPartitionerShellGetRandom - get the flag to use a random partition

    223:   Collective on PetscPartitioner

    225:   Input Parameter:
    226: . part   - The PetscPartitioner

    228:   Output Parameter:
    229: . random - The flag to use a random partition

    231:   Level: intermediate

    233: .seealso PetscPartitionerShellSetRandom(), PetscPartitionerCreate()
    234: @*/
    235: PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random)
    236: {
    237:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;

    242:   *random = p->random;
    243:   return(0);
    244: }