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:   PetscFunctionBegin;
     14:   PetscCall(PetscSectionDestroy(&p->section));
     15:   PetscCall(ISDestroy(&p->partition));
     16:   PetscFunctionReturn(PETSC_SUCCESS);
     17: }

     19: static PetscErrorCode PetscPartitionerDestroy_Shell(PetscPartitioner part)
     20: {
     21:   PetscFunctionBegin;
     22:   PetscCall(PetscPartitionerReset_Shell(part));
     23:   PetscCall(PetscFree(part->data));
     24:   PetscFunctionReturn(PETSC_SUCCESS);
     25: }

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

     31:   PetscFunctionBegin;
     32:   if (p->random) {
     33:     PetscCall(PetscViewerASCIIPushTab(viewer));
     34:     PetscCall(PetscViewerASCIIPrintf(viewer, "using random partition\n"));
     35:     PetscCall(PetscViewerASCIIPopTab(viewer));
     36:   }
     37:   PetscFunctionReturn(PETSC_SUCCESS);
     38: }

     40: static PetscErrorCode PetscPartitionerView_Shell(PetscPartitioner part, PetscViewer viewer)
     41: {
     42:   PetscBool isascii;

     44:   PetscFunctionBegin;
     47:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
     48:   if (isascii) PetscCall(PetscPartitionerView_Shell_ASCII(part, viewer));
     49:   PetscFunctionReturn(PETSC_SUCCESS);
     50: }

     52: static PetscErrorCode PetscPartitionerSetFromOptions_Shell(PetscPartitioner part, PetscOptionItems PetscOptionsObject)
     53: {
     54:   PetscInt    sizes[16], points[1024];
     55:   PetscInt    Npart = 16, Npoints = 1024;
     56:   PetscBool   random = PETSC_FALSE, set, flgSizes, flgPoints;
     57:   PetscMPIInt rank;

     59:   PetscFunctionBegin;
     60:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank));
     61:   PetscOptionsHeadBegin(PetscOptionsObject, "PetscPartitioner Shell Options");
     62:   PetscCall(PetscPartitionerShellGetRandom(part, &random));
     63:   PetscCall(PetscOptionsBool("-petscpartitioner_shell_random", "Use a random partition", "PetscPartitionerView", PETSC_FALSE, &random, &set));
     64:   if (set) PetscCall(PetscPartitionerShellSetRandom(part, random));
     65:   PetscCall(PetscOptionsIntArray("-petscpartitioner_shell_sizes", "The size of each partition on rank 0", "PetscPartitionerShellSetPartition", sizes, &Npart, &flgSizes));
     66:   PetscCall(PetscOptionsIntArray("-petscpartitioner_shell_points", "The points in each partition on rank 0", "PetscPartitionerShellSetPartition", points, &Npoints, &flgPoints));
     67:   PetscCheck(!(flgSizes ^ flgPoints), PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Must specify both the partition sizes and points");
     68:   if (flgSizes) {
     69:     PetscInt Np = 0;

     71:     for (PetscInt i = 0; i < Npart; ++i) Np += sizes[i];
     72:     PetscCheck(Np == Npoints, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Number of input points %" PetscInt_FMT " != %" PetscInt_FMT " sum of partition sizes", Npoints, Np);
     73:     if (!rank) PetscCall(PetscPartitionerShellSetPartition(part, Npart, sizes, points));
     74:     else {
     75:       PetscCall(PetscArrayzero(sizes, Npart));
     76:       PetscCall(PetscPartitionerShellSetPartition(part, Npart, sizes, points));
     77:     }
     78:   }
     79:   PetscOptionsHeadEnd();
     80:   PetscFunctionReturn(PETSC_SUCCESS);
     81: }

     83: static PetscErrorCode PetscPartitionerPartition_Shell(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *partition)
     84: {
     85:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
     86:   PetscInt                np;

     88:   PetscFunctionBegin;
     89:   if (p->random) {
     90:     PetscRandom r;
     91:     PetscInt   *sizes, *points, v, p;
     92:     PetscMPIInt rank;

     94:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)part), &rank));
     95:     PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &r));
     96:     PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)nparts));
     97:     PetscCall(PetscRandomSetFromOptions(r));
     98:     PetscCall(PetscCalloc2(nparts, &sizes, numVertices, &points));
     99:     for (v = 0; v < numVertices; ++v) points[v] = v;
    100:     for (p = 0; p < nparts; ++p) sizes[p] = numVertices / nparts + (PetscInt)(p < numVertices % nparts);
    101:     for (v = numVertices - 1; v > 0; --v) {
    102:       PetscReal val;
    103:       PetscInt  w, tmp;

    105:       PetscCall(PetscRandomSetInterval(r, 0.0, (PetscScalar)(v + 1)));
    106:       PetscCall(PetscRandomGetValueReal(r, &val));
    107:       w         = PetscFloorReal(val);
    108:       tmp       = points[v];
    109:       points[v] = points[w];
    110:       points[w] = tmp;
    111:     }
    112:     PetscCall(PetscRandomDestroy(&r));
    113:     PetscCall(PetscPartitionerShellSetPartition(part, nparts, sizes, points));
    114:     PetscCall(PetscFree2(sizes, points));
    115:   }
    116:   PetscCheck(p->section, PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_WRONG, "Shell partitioner information not provided. Please call PetscPartitionerShellSetPartition()");
    117:   PetscCall(PetscSectionGetChart(p->section, NULL, &np));
    118:   PetscCheck(nparts == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of requested partitions %" PetscInt_FMT " != configured partitions %" PetscInt_FMT, nparts, np);
    119:   PetscCall(ISGetLocalSize(p->partition, &np));
    120:   PetscCheck(numVertices == np, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Number of input vertices %" PetscInt_FMT " != configured vertices %" PetscInt_FMT, numVertices, np);
    121:   PetscCall(PetscSectionCopy(p->section, partSection));
    122:   *partition = p->partition;
    123:   PetscCall(PetscObjectReference((PetscObject)p->partition));
    124:   PetscFunctionReturn(PETSC_SUCCESS);
    125: }

    127: static PetscErrorCode PetscPartitionerInitialize_Shell(PetscPartitioner part)
    128: {
    129:   PetscFunctionBegin;
    130:   part->noGraph             = PETSC_TRUE; /* PetscPartitionerShell cannot overload the partition call, so it is safe for now */
    131:   part->ops->view           = PetscPartitionerView_Shell;
    132:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_Shell;
    133:   part->ops->reset          = PetscPartitionerReset_Shell;
    134:   part->ops->destroy        = PetscPartitionerDestroy_Shell;
    135:   part->ops->partition      = PetscPartitionerPartition_Shell;
    136:   PetscFunctionReturn(PETSC_SUCCESS);
    137: }

    139: /*MC
    140:   PETSCPARTITIONERSHELL = "shell" - A PetscPartitioner object

    142:   Level: intermediate

    144:   Options Database Keys:
    145: .  -petscpartitioner_shell_random - Use a random partition

    147: .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
    148: M*/

    150: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_Shell(PetscPartitioner part)
    151: {
    152:   PetscPartitioner_Shell *p;

    154:   PetscFunctionBegin;
    156:   PetscCall(PetscNew(&p));
    157:   part->data = p;

    159:   PetscCall(PetscPartitionerInitialize_Shell(part));
    160:   p->random = PETSC_FALSE;
    161:   PetscFunctionReturn(PETSC_SUCCESS);
    162: }

    164: /*@
    165:   PetscPartitionerShellSetPartition - Set an artificial partition for a mesh

    167:   Collective

    169:   Input Parameters:
    170: + part   - The `PetscPartitioner`
    171: . size   - The number of partitions
    172: . sizes  - array of length size (or `NULL`) providing the number of points in each partition
    173: - 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.)

    175:   Level: developer

    177:   Note:
    178:   It is safe to free the sizes and points arrays after use in this routine.

    180: .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()`
    181: @*/
    182: PetscErrorCode PetscPartitionerShellSetPartition(PetscPartitioner part, PetscInt size, const PetscInt sizes[], const PetscInt points[])
    183: {
    184:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;
    185:   PetscInt                proc, numPoints;

    187:   PetscFunctionBegin;
    189:   if (sizes) PetscAssertPointer(sizes, 3);
    190:   if (points) PetscAssertPointer(points, 4);
    191:   PetscCall(PetscSectionDestroy(&p->section));
    192:   PetscCall(ISDestroy(&p->partition));
    193:   PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)part), &p->section));
    194:   PetscCall(PetscSectionSetChart(p->section, 0, size));
    195:   if (sizes) {
    196:     for (proc = 0; proc < size; ++proc) PetscCall(PetscSectionSetDof(p->section, proc, sizes[proc]));
    197:   }
    198:   PetscCall(PetscSectionSetUp(p->section));
    199:   PetscCall(PetscSectionGetStorageSize(p->section, &numPoints));
    200:   PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)part), numPoints, points, PETSC_COPY_VALUES, &p->partition));
    201:   PetscFunctionReturn(PETSC_SUCCESS);
    202: }

    204: /*@
    205:   PetscPartitionerShellSetRandom - Set the flag to use a random partition

    207:   Collective

    209:   Input Parameters:
    210: + part   - The `PetscPartitioner`
    211: - random - The flag to use a random partition

    213:   Level: intermediate

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

    221:   PetscFunctionBegin;
    223:   p->random = random;
    224:   PetscFunctionReturn(PETSC_SUCCESS);
    225: }

    227: /*@
    228:   PetscPartitionerShellGetRandom - get the flag to use a random partition

    230:   Collective

    232:   Input Parameter:
    233: . part - The `PetscPartitioner`

    235:   Output Parameter:
    236: . random - The flag to use a random partition

    238:   Level: intermediate

    240: .seealso: `PetscPartitionerShellSetRandom()`, `PetscPartitionerCreate()`
    241: @*/
    242: PetscErrorCode PetscPartitionerShellGetRandom(PetscPartitioner part, PetscBool *random)
    243: {
    244:   PetscPartitioner_Shell *p = (PetscPartitioner_Shell *)part->data;

    246:   PetscFunctionBegin;
    248:   PetscAssertPointer(random, 2);
    249:   *random = p->random;
    250:   PetscFunctionReturn(PETSC_SUCCESS);
    251: }