Actual source code: stride.c

      1: /*
      2:        Index sets of evenly space integers, defined by a
      3:     start, stride and length.
      4: */
      5: #include <petsc/private/isimpl.h>
      6: #include <petscviewer.h>

      8: typedef struct {
      9:   PetscInt first, step;
     10: } IS_Stride;

     12: static PetscErrorCode ISCopy_Stride(IS is, IS isy)
     13: {
     14:   IS_Stride *is_stride = (IS_Stride *)is->data, *isy_stride = (IS_Stride *)isy->data;

     16:   PetscFunctionBegin;
     17:   PetscCall(PetscMemcpy(isy_stride, is_stride, sizeof(IS_Stride)));
     18:   PetscFunctionReturn(PETSC_SUCCESS);
     19: }

     21: static PetscErrorCode ISShift_Stride(IS is, PetscInt shift, IS isy)
     22: {
     23:   IS_Stride *is_stride = (IS_Stride *)is->data, *isy_stride = (IS_Stride *)isy->data;

     25:   PetscFunctionBegin;
     26:   isy_stride->first = is_stride->first + shift;
     27:   isy_stride->step  = is_stride->step;
     28:   PetscFunctionReturn(PETSC_SUCCESS);
     29: }

     31: static PetscErrorCode ISDuplicate_Stride(IS is, IS *newIS)
     32: {
     33:   IS_Stride *sub = (IS_Stride *)is->data;

     35:   PetscFunctionBegin;
     36:   PetscCall(ISCreateStride(PetscObjectComm((PetscObject)is), is->map->n, sub->first, sub->step, newIS));
     37:   PetscFunctionReturn(PETSC_SUCCESS);
     38: }

     40: static PetscErrorCode ISInvertPermutation_Stride(IS is, PetscInt nlocal, IS *perm)
     41: {
     42:   PetscBool isident, samelocal = (PetscBool)(nlocal == PETSC_DECIDE);

     44:   PetscFunctionBegin;
     45:   PetscCall(ISGetInfo(is, IS_IDENTITY, IS_GLOBAL, PETSC_TRUE, &isident));
     46:   if (isident && nlocal != PETSC_DECIDE) PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &samelocal, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)is)));
     47:   if (isident) {
     48:     PetscInt start = is->map->rstart, n = is->map->n;

     50:     if (!samelocal) {
     51:       n     = nlocal;
     52:       start = 0;

     54:       PetscCallMPI(MPI_Exscan(&nlocal, &start, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)is)));
     55:     }
     56:     PetscCall(ISCreateStride(PetscObjectComm((PetscObject)is), n, start, 1, perm));
     57:   } else {
     58:     IS              tmp;
     59:     const PetscInt *indices, n = is->map->n;

     61:     PetscCall(ISGetIndices(is, &indices));
     62:     PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is), n, indices, PETSC_COPY_VALUES, &tmp));
     63:     PetscCall(ISSetPermutation(tmp));
     64:     PetscCall(ISRestoreIndices(is, &indices));
     65:     PetscCall(ISInvertPermutation(tmp, nlocal, perm));
     66:     PetscCall(ISDestroy(&tmp));
     67:   }
     68:   PetscFunctionReturn(PETSC_SUCCESS);
     69: }

     71: /*@
     72:   ISStrideGetInfo - Returns the first index in a stride index set and the stride width from an `IS` of `ISType` `ISSTRIDE`

     74:   Not Collective

     76:   Input Parameter:
     77: . is - the index set

     79:   Output Parameters:
     80: + first - the first index
     81: - step  - the stride width

     83:   Level: intermediate

     85: .seealso: [](sec_scatter), `IS`, `ISCreateStride()`, `ISGetSize()`, `ISSTRIDE`
     86: @*/
     87: PetscErrorCode ISStrideGetInfo(IS is, PetscInt *first, PetscInt *step)
     88: {
     89:   IS_Stride *sub;
     90:   PetscBool  flg;

     92:   PetscFunctionBegin;
     94:   if (first) PetscAssertPointer(first, 2);
     95:   if (step) PetscAssertPointer(step, 3);
     96:   PetscCall(PetscObjectTypeCompare((PetscObject)is, ISSTRIDE, &flg));
     97:   PetscCheck(flg, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_WRONG, "IS must be of type ISSTRIDE");

     99:   sub = (IS_Stride *)is->data;
    100:   if (first) *first = sub->first;
    101:   if (step) *step = sub->step;
    102:   PetscFunctionReturn(PETSC_SUCCESS);
    103: }

    105: static PetscErrorCode ISDestroy_Stride(IS is)
    106: {
    107:   PetscFunctionBegin;
    108:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", NULL));
    109:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISShift_C", NULL));
    110:   PetscCall(PetscFree(is->data));
    111:   PetscFunctionReturn(PETSC_SUCCESS);
    112: }

    114: static PetscErrorCode ISToGeneral_Stride(IS inis)
    115: {
    116:   const PetscInt *idx;
    117:   PetscInt        n;

    119:   PetscFunctionBegin;
    120:   PetscCall(ISGetLocalSize(inis, &n));
    121:   PetscCall(ISGetIndices(inis, &idx));
    122:   PetscCall(ISSetType(inis, ISGENERAL));
    123:   PetscCall(ISGeneralSetIndices(inis, n, idx, PETSC_OWN_POINTER));
    124:   PetscFunctionReturn(PETSC_SUCCESS);
    125: }

    127: static PetscErrorCode ISLocate_Stride(IS is, PetscInt key, PetscInt *location)
    128: {
    129:   IS_Stride *sub = (IS_Stride *)is->data;
    130:   PetscInt   rem, step;

    132:   PetscFunctionBegin;
    133:   *location = -1;
    134:   step      = sub->step;
    135:   key -= sub->first;
    136:   rem = key / step;
    137:   if ((rem < is->map->n) && !(key % step)) *location = rem;
    138:   PetscFunctionReturn(PETSC_SUCCESS);
    139: }

    141: /*
    142:      Returns a legitimate index memory even if
    143:    the stride index set is empty.
    144: */
    145: static PetscErrorCode ISGetIndices_Stride(IS is, const PetscInt *idx[])
    146: {
    147:   IS_Stride *sub = (IS_Stride *)is->data;
    148:   PetscInt   i, **dx = (PetscInt **)idx;

    150:   PetscFunctionBegin;
    151:   PetscCall(PetscMalloc1(is->map->n, (PetscInt **)idx));
    152:   if (is->map->n) {
    153:     (*dx)[0] = sub->first;
    154:     for (i = 1; i < is->map->n; i++) (*dx)[i] = (*dx)[i - 1] + sub->step;
    155:   }
    156:   PetscFunctionReturn(PETSC_SUCCESS);
    157: }

    159: static PetscErrorCode ISRestoreIndices_Stride(IS in, const PetscInt *idx[])
    160: {
    161:   PetscFunctionBegin;
    162:   PetscCall(PetscFree(*(void **)idx));
    163:   PetscFunctionReturn(PETSC_SUCCESS);
    164: }

    166: static PetscErrorCode ISView_Stride(IS is, PetscViewer viewer)
    167: {
    168:   IS_Stride        *sub = (IS_Stride *)is->data;
    169:   PetscInt          i, n = is->map->n;
    170:   PetscMPIInt       rank, size;
    171:   PetscBool         isascii, ibinary;
    172:   PetscViewerFormat fmt;

    174:   PetscFunctionBegin;
    175:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    176:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &ibinary));
    177:   if (isascii) {
    178:     PetscBool matl, isperm;

    180:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)is), &rank));
    181:     PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)is), &size));
    182:     PetscCall(PetscViewerGetFormat(viewer, &fmt));
    183:     matl = (PetscBool)(fmt == PETSC_VIEWER_ASCII_MATLAB);
    184:     PetscCall(ISGetInfo(is, IS_PERMUTATION, IS_GLOBAL, PETSC_FALSE, &isperm));
    185:     if (isperm && !matl) PetscCall(PetscViewerASCIIPrintf(viewer, "Index set is permutation\n"));
    186:     if (size == 1) {
    187:       if (matl) {
    188:         const char *name;

    190:         PetscCall(PetscObjectGetName((PetscObject)is, &name));
    191:         PetscCall(PetscViewerASCIIPrintf(viewer, "%s = [%" PetscInt_FMT " : %" PetscInt_FMT " : %" PetscInt_FMT "];\n", name, sub->first + 1, sub->step, sub->first + sub->step * (n - 1) + 1));
    192:       } else {
    193:         PetscCall(PetscViewerASCIIPrintf(viewer, "Number of indices in (stride) set %" PetscInt_FMT "\n", n));
    194:         for (i = 0; i < n; i++) PetscCall(PetscViewerASCIIPrintf(viewer, "%" PetscInt_FMT " %" PetscInt_FMT "\n", i, sub->first + i * sub->step));
    195:       }
    196:       PetscCall(PetscViewerFlush(viewer));
    197:     } else {
    198:       PetscCall(PetscViewerASCIIPushSynchronized(viewer));
    199:       if (matl) {
    200:         const char *name;

    202:         PetscCall(PetscObjectGetName((PetscObject)is, &name));
    203:         PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "%s_%d = [%" PetscInt_FMT " : %" PetscInt_FMT " : %" PetscInt_FMT "];\n", name, rank, sub->first + 1, sub->step, sub->first + sub->step * (n - 1) + 1));
    204:       } else {
    205:         PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] Number of indices in (stride) set %" PetscInt_FMT "\n", rank, n));
    206:         for (i = 0; i < n; i++) PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] %" PetscInt_FMT " %" PetscInt_FMT "\n", rank, i, sub->first + i * sub->step));
    207:       }
    208:       PetscCall(PetscViewerFlush(viewer));
    209:       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
    210:     }
    211:   } else if (ibinary) PetscCall(ISView_Binary(is, viewer));
    212:   PetscFunctionReturn(PETSC_SUCCESS);
    213: }

    215: static PetscErrorCode ISSort_Stride(IS is)
    216: {
    217:   IS_Stride *sub = (IS_Stride *)is->data;

    219:   PetscFunctionBegin;
    220:   if (sub->step >= 0) PetscFunctionReturn(PETSC_SUCCESS);
    221:   sub->first += (is->map->n - 1) * sub->step;
    222:   sub->step *= -1;
    223:   PetscFunctionReturn(PETSC_SUCCESS);
    224: }

    226: static PetscErrorCode ISSorted_Stride(IS is, PetscBool *flg)
    227: {
    228:   IS_Stride *sub = (IS_Stride *)is->data;

    230:   PetscFunctionBegin;
    231:   if (sub->step >= 0) *flg = PETSC_TRUE;
    232:   else *flg = PETSC_FALSE;
    233:   PetscFunctionReturn(PETSC_SUCCESS);
    234: }

    236: static PetscErrorCode ISUniqueLocal_Stride(IS is, PetscBool *flg)
    237: {
    238:   IS_Stride *sub = (IS_Stride *)is->data;

    240:   PetscFunctionBegin;
    241:   if (!is->map->n || sub->step != 0) *flg = PETSC_TRUE;
    242:   else *flg = PETSC_FALSE;
    243:   PetscFunctionReturn(PETSC_SUCCESS);
    244: }

    246: static PetscErrorCode ISPermutationLocal_Stride(IS is, PetscBool *flg)
    247: {
    248:   IS_Stride *sub = (IS_Stride *)is->data;

    250:   PetscFunctionBegin;
    251:   if (!is->map->n || (PetscAbsInt(sub->step) == 1 && is->min == 0)) *flg = PETSC_TRUE;
    252:   else *flg = PETSC_FALSE;
    253:   PetscFunctionReturn(PETSC_SUCCESS);
    254: }

    256: static PetscErrorCode ISIntervalLocal_Stride(IS is, PetscBool *flg)
    257: {
    258:   IS_Stride *sub = (IS_Stride *)is->data;

    260:   PetscFunctionBegin;
    261:   if (!is->map->n || sub->step == 1) *flg = PETSC_TRUE;
    262:   else *flg = PETSC_FALSE;
    263:   PetscFunctionReturn(PETSC_SUCCESS);
    264: }

    266: static PetscErrorCode ISOnComm_Stride(IS is, MPI_Comm comm, PetscCopyMode mode, IS *newis)
    267: {
    268:   IS_Stride *sub = (IS_Stride *)is->data;

    270:   PetscFunctionBegin;
    271:   PetscCall(ISCreateStride(comm, is->map->n, sub->first, sub->step, newis));
    272:   PetscFunctionReturn(PETSC_SUCCESS);
    273: }

    275: static PetscErrorCode ISSetBlockSize_Stride(IS is, PetscInt bs)
    276: {
    277:   IS_Stride *sub = (IS_Stride *)is->data;

    279:   PetscFunctionBegin;
    280:   PetscCheck(sub->step == 1 || bs == 1, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_SIZ, "ISSTRIDE has stride %" PetscInt_FMT ", cannot be blocked of size %" PetscInt_FMT, sub->step, bs);
    281:   PetscCall(PetscLayoutSetBlockSize(is->map, bs));
    282:   PetscFunctionReturn(PETSC_SUCCESS);
    283: }

    285: static PetscErrorCode ISContiguousLocal_Stride(IS is, PetscInt gstart, PetscInt gend, PetscInt *start, PetscBool *contig)
    286: {
    287:   IS_Stride *sub = (IS_Stride *)is->data;

    289:   PetscFunctionBegin;
    290:   if (sub->step == 1 && sub->first >= gstart && sub->first + is->map->n <= gend) {
    291:     *start  = sub->first - gstart;
    292:     *contig = PETSC_TRUE;
    293:   } else {
    294:     *start  = -1;
    295:     *contig = PETSC_FALSE;
    296:   }
    297:   PetscFunctionReturn(PETSC_SUCCESS);
    298: }

    300: // clang-format off
    301: static const struct _ISOps myops = {
    302:   PetscDesignatedInitializer(getindices, ISGetIndices_Stride),
    303:   PetscDesignatedInitializer(restoreindices, ISRestoreIndices_Stride),
    304:   PetscDesignatedInitializer(invertpermutation, ISInvertPermutation_Stride),
    305:   PetscDesignatedInitializer(sort, ISSort_Stride),
    306:   PetscDesignatedInitializer(sortremovedups, ISSort_Stride),
    307:   PetscDesignatedInitializer(sorted, ISSorted_Stride),
    308:   PetscDesignatedInitializer(duplicate, ISDuplicate_Stride),
    309:   PetscDesignatedInitializer(destroy, ISDestroy_Stride),
    310:   PetscDesignatedInitializer(view, ISView_Stride),
    311:   PetscDesignatedInitializer(load, ISLoad_Default),
    312:   PetscDesignatedInitializer(copy, ISCopy_Stride),
    313:   PetscDesignatedInitializer(togeneral, ISToGeneral_Stride),
    314:   PetscDesignatedInitializer(oncomm, ISOnComm_Stride),
    315:   PetscDesignatedInitializer(setblocksize, ISSetBlockSize_Stride),
    316:   PetscDesignatedInitializer(contiguous, ISContiguousLocal_Stride),
    317:   PetscDesignatedInitializer(locate, ISLocate_Stride),
    318:   PetscDesignatedInitializer(sortedlocal, ISSorted_Stride),
    319:   PetscDesignatedInitializer(sortedglobal, NULL),
    320:   PetscDesignatedInitializer(uniquelocal, ISUniqueLocal_Stride),
    321:   PetscDesignatedInitializer(uniqueglobal, NULL),
    322:   PetscDesignatedInitializer(permlocal, ISPermutationLocal_Stride),
    323:   PetscDesignatedInitializer(permglobal, NULL),
    324:   PetscDesignatedInitializer(intervallocal, ISIntervalLocal_Stride),
    325:   PetscDesignatedInitializer(intervalglobal, NULL)
    326: };
    327: // clang-format on

    329: /*@
    330:   ISStrideSetStride - Sets the stride information for a stride index set.

    332:   Logically Collective

    334:   Input Parameters:
    335: + is    - the index set
    336: . n     - the length of the locally owned portion of the index set
    337: . first - the first element of the locally owned portion of the index set
    338: - step  - the change to the next index

    340:   Level: beginner

    342:   Note:
    343:   `ISCreateStride()` can be used to create an `ISSTRIDE` and set its stride in one function call

    345: .seealso: [](sec_scatter), `IS`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`, `ISCreateStride()`, `ISStrideGetInfo()`
    346: @*/
    347: PetscErrorCode ISStrideSetStride(IS is, PetscInt n, PetscInt first, PetscInt step)
    348: {
    349:   PetscFunctionBegin;
    350:   PetscCheck(n >= 0, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_OUTOFRANGE, "Negative length %" PetscInt_FMT " not valid", n);
    351:   PetscCall(ISClearInfoCache(is, PETSC_FALSE));
    352:   PetscUseMethod(is, "ISStrideSetStride_C", (IS, PetscInt, PetscInt, PetscInt), (is, n, first, step));
    353:   PetscFunctionReturn(PETSC_SUCCESS);
    354: }

    356: static PetscErrorCode ISStrideSetStride_Stride(IS is, PetscInt n, PetscInt first, PetscInt step)
    357: {
    358:   PetscInt    min, max;
    359:   IS_Stride  *sub = (IS_Stride *)is->data;
    360:   PetscLayout map;

    362:   PetscFunctionBegin;
    363:   PetscCall(PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)is), n, is->map->N, is->map->bs, &map));
    364:   PetscCall(PetscLayoutDestroy(&is->map));
    365:   is->map = map;

    367:   sub->first = first;
    368:   sub->step  = step;
    369:   if (step > 0) {
    370:     min = first;
    371:     max = first + step * (n - 1);
    372:   } else {
    373:     max = first;
    374:     min = first + step * (n - 1);
    375:   }

    377:   is->min  = n > 0 ? min : PETSC_INT_MAX;
    378:   is->max  = n > 0 ? max : PETSC_INT_MIN;
    379:   is->data = (void *)sub;
    380:   PetscFunctionReturn(PETSC_SUCCESS);
    381: }

    383: /*@
    384:   ISCreateStride - Creates a data structure for an index set containing a list of evenly spaced integers.

    386:   Collective

    388:   Input Parameters:
    389: + comm  - the MPI communicator
    390: . n     - the length of the locally owned portion of the index set
    391: . first - the first element of the locally owned portion of the index set
    392: - step  - the change to the next index

    394:   Output Parameter:
    395: . is - the new index set

    397:   Level: beginner

    399:   Notes:
    400:   `ISStrideSetStride()` may be used to set the stride of an `ISSTRIDE` that already exists

    402:   When the communicator is not `MPI_COMM_SELF`, the operations on `IS` are NOT
    403:   conceptually the same as `MPI_Group` operations. The `IS` are the
    404:   distributed sets of indices and thus certain operations on them are collective.

    406: .seealso: [](sec_scatter), `IS`, `ISStrideSetStride()`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`
    407: @*/
    408: PetscErrorCode ISCreateStride(MPI_Comm comm, PetscInt n, PetscInt first, PetscInt step, IS *is)
    409: {
    410:   PetscFunctionBegin;
    411:   PetscCall(ISCreate(comm, is));
    412:   PetscCall(ISSetType(*is, ISSTRIDE));
    413:   PetscCall(ISStrideSetStride(*is, n, first, step));
    414:   PetscFunctionReturn(PETSC_SUCCESS);
    415: }

    417: PETSC_INTERN PetscErrorCode ISCreate_Stride(IS is)
    418: {
    419:   IS_Stride *sub;

    421:   PetscFunctionBegin;
    422:   PetscCall(PetscNew(&sub));
    423:   is->data   = (void *)sub;
    424:   is->ops[0] = myops;
    425:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", ISStrideSetStride_Stride));
    426:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISShift_C", ISShift_Stride));
    427:   PetscFunctionReturn(PETSC_SUCCESS);
    428: }