Date: 2022-03-31T04:47:55+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;

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

     18: static PetscErrorCode PetscPartitionerDestroy_Shell(PetscPartitioner part)
     19: {
     20:   PetscPartitionerReset_Shell(part);
     21:   PetscFree(part->data);
     22:   return 0;
     23: }

     25: static PetscErrorCode PetscPartitionerView_Shell_ASCII(PetscPartitioner part, PetscViewer viewer)
     26: {
     27:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;

     29:   if (p->random) {
     30:     PetscViewerASCIIPushTab(viewer);
     31:     PetscViewerASCIIPrintf(viewer, "using random partition\n");
     32:     PetscViewerASCIIPopTab(viewer);
     33:   }
     34:   return 0;
     35: }

     37: static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer)
     38: {
     39:   PetscBool      iascii;

     43:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
     44:   if (iascii) PetscPartitionerView_Shell_ASCII(part, viewer);
     45:   return 0;
     46: }

     48: static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscOptionItems *PetscOptionsObject, PetscPartitioner part)
     49: {
     50:   PetscBool      random = PETSC_FALSE, set;

     52:   PetscOptionsHead(PetscOptionsObject, "PetscPartitioner Shell Options");
     53:   PetscPartitionerShellGetRandom(part, &random);
     54:   PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set);
     55:   if (set) PetscPartitionerShellSetRandom(part, random);
     56:   PetscOptionsTail();
     57:   return 0;
     58: }

     60: static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection targetSection, PetscSection partSection, IS *partition)
     61: {
     62:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;
     63:   PetscInt                np;

     65:   if (p->random) {
     66:     PetscRandom r;
     67:     PetscInt   *sizes, *points, v, p;
     68:     PetscMPIInt rank;

     70:     MPI_Comm_rank(PetscObjectComm((PetscObject) part), &rank);
     71:     PetscRandomCreate(PETSC_COMM_SELF, &r);
     72:     PetscRandomSetInterval(r, 0.0, (PetscScalar) nparts);
     73:     PetscRandomSetFromOptions(r);
     74:     PetscCalloc2(nparts, &sizes, numVertices, &points);
     75:     for (v = 0; v < numVertices; ++v) {points[v] = v;}
     76:     for (p = 0; p < nparts; ++p) {sizes[p] = numVertices/nparts + (PetscInt) (p < numVertices % nparts);}
     77:     for (v = numVertices-1; v > 0; --v) {
     78:       PetscReal val;
     79:       PetscInt  w, tmp;

     81:       PetscRandomSetInterval(r, 0.0, (PetscScalar) (v+1));
     82:       PetscRandomGetValueReal(r, &val);
     83:       w    = PetscFloorReal(val);
     84:       tmp       = points[v];
     85:       points[v] = points[w];
     86:       points[w] = tmp;
     87:     }
     88:     PetscRandomDestroy(&r);
     89:     PetscPartitionerShellSetPartition(part, nparts, sizes, points);
     90:     PetscFree2(sizes, points);
     91:   }
     93:   PetscSectionGetChart(p->section, NULL, &np);
     95:   ISGetLocalSize(p->partition, &np);
     97:   PetscSectionCopy(p->section, partSection);
     98:   *partition = p->partition;
     99:   PetscObjectReference((PetscObject) p->partition);
    100:   return 0;
    101: }

    103: static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part)
    104: {
    105:   part->noGraph             = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */
    106:   part->ops->view           = PetscPartitionerView_Shell;
    107:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell;
    108:   part->ops->reset          = PetscPartitionerReset_Shell;
    109:   part->ops->destroy        = PetscPartitionerDestroy_Shell;
    110:   part->ops->partition      = PetscPartitionerPartition_Shell;
    111:   return 0;
    112: }

    114: /*MC
    115:   PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object

    117:   Level: intermediate

    119:   Options Database Keys:
    120: .  -petscpartitioner_shell_random - Use a random partition

    122: .seealso: PetscPartitionerType, PetscPartitionerCreate(), PetscPartitionerSetType()
    123: M*/

    125: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part)
    126: {
    127:   PetscPartitioner_Shell *p;

    130:   PetscNewLog(part, &p);
    131:   part->data = p;

    133:   PetscPartitionerInitialize_Shell(part);
    134:   p->random = PETSC_FALSE;
    135:   return 0;
    136: }

    138: /*@C
    139:   PetscPartitionerShellSetPartition - Set an artifical partition for a mesh

    141:   Collective on PetscPartitioner

    143:   Input Parameters:
    144: + part   - The PetscPartitioner
    145: . size   - The number of partitions
    146: . sizes  - array of length size (or NULL) providing the number of points in each partition
    147: - 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.)

    149:   Level: developer

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

    154: .seealso DMPlexDistribute(), PetscPartitionerCreate()
    155: @*/
    156: PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[])
    157: {
    158:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;
    159:   PetscInt                proc, numPoints;

    164:   PetscSectionDestroy(&p->section);
    165:   ISDestroy(&p->partition);
    166:   PetscSectionCreate(PetscObjectComm((PetscObject) part), &p->section);
    167:   PetscSectionSetChart(p->section, 0, size);
    168:   if (sizes) {
    169:     for (proc = 0; proc < size; ++proc) {
    170:       PetscSectionSetDof(p->section, proc, sizes[proc]);
    171:     }
    172:   }
    173:   PetscSectionSetUp(p->section);
    174:   PetscSectionGetStorageSize(p->section, &numPoints);
    175:   ISCreateGeneral(PetscObjectComm((PetscObject) part), numPoints, points, PETSC_COPY_VALUES, &p->partition);
    176:   return 0;
    177: }

    179: /*@
    180:   PetscPartitionerShellSetRandom - Set the flag to use a random partition

    182:   Collective on PetscPartitioner

    184:   Input Parameters:
    185: + part   - The PetscPartitioner
    186: - random - The flag to use a random partition

    188:   Level: intermediate

    190: .seealso PetscPartitionerShellGetRandom(), PetscPartitionerCreate()
    191: @*/
    192: PetscErrorCode PetscPartitionerShellSetRandom(PetscPartitioner part, PetscBool random)
    193: {
    194:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;

    197:   p->random = random;
    198:   return 0;
    199: }

    201: /*@
    202:   PetscPartitionerShellGetRandom - get the flag to use a random partition

    204:   Collective on PetscPartitioner

    206:   Input Parameter:
    207: . part   - The PetscPartitioner

    209:   Output Parameter:
    210: . random - The flag to use a random partition

    212:   Level: intermediate

    214: .seealso PetscPartitionerShellSetRandom(), PetscPartitionerCreate()
    215: @*/
    216: PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random)
    217: {
    218:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *) part->data;

    222:   *random = p->random;
    223:   return 0;
    224: }