Actual source code: stride.c


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

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

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

 17:   PetscMemcpy(isy_stride, is_stride, sizeof(IS_Stride));
 18:   return 0;
 19: }

 21: 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:   isy_stride->first = is_stride->first + shift;
 26:   isy_stride->step  = is_stride->step;
 27:   return 0;
 28: }

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

 34:   ISCreateStride(PetscObjectComm((PetscObject)is), is->map->n, sub->first, sub->step, newIS);
 35:   return 0;
 36: }

 38: PetscErrorCode ISInvertPermutation_Stride(IS is, PetscInt nlocal, IS *perm)
 39: {
 40:   PetscBool isident;

 42:   ISGetInfo(is, IS_IDENTITY, IS_GLOBAL, PETSC_TRUE, &isident);
 43:   if (isident) {
 44:     PetscInt rStart, rEnd;

 46:     PetscLayoutGetRange(is->map, &rStart, &rEnd);
 47:     ISCreateStride(PETSC_COMM_SELF, PetscMax(rEnd - rStart, 0), rStart, 1, perm);
 48:   } else {
 49:     IS              tmp;
 50:     const PetscInt *indices, n = is->map->n;

 52:     ISGetIndices(is, &indices);
 53:     ISCreateGeneral(PetscObjectComm((PetscObject)is), n, indices, PETSC_COPY_VALUES, &tmp);
 54:     ISSetPermutation(tmp);
 55:     ISRestoreIndices(is, &indices);
 56:     ISInvertPermutation(tmp, nlocal, perm);
 57:     ISDestroy(&tmp);
 58:   }
 59:   return 0;
 60: }

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

 65:    Not Collective

 67:    Input Parameter:
 68: .  is - the index set

 70:    Output Parameters:
 71: +  first - the first index
 72: -  step - the stride width

 74:    Level: intermediate

 76: .seealso: [](sec_scatter), `IS`, `ISCreateStride()`, `ISGetSize()`, `ISSTRIDE`
 77: @*/
 78: PetscErrorCode ISStrideGetInfo(IS is, PetscInt *first, PetscInt *step)
 79: {
 80:   IS_Stride *sub;
 81:   PetscBool  flg;

 86:   PetscObjectTypeCompare((PetscObject)is, ISSTRIDE, &flg);

 89:   sub = (IS_Stride *)is->data;
 90:   if (first) *first = sub->first;
 91:   if (step) *step = sub->step;
 92:   return 0;
 93: }

 95: PetscErrorCode ISDestroy_Stride(IS is)
 96: {
 97:   PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", NULL);
 98:   PetscObjectComposeFunction((PetscObject)is, "ISShift_C", NULL);
 99:   PetscFree(is->data);
100:   return 0;
101: }

103: PetscErrorCode ISToGeneral_Stride(IS inis)
104: {
105:   const PetscInt *idx;
106:   PetscInt        n;

108:   ISGetLocalSize(inis, &n);
109:   ISGetIndices(inis, &idx);
110:   ISSetType(inis, ISGENERAL);
111:   ISGeneralSetIndices(inis, n, idx, PETSC_OWN_POINTER);
112:   return 0;
113: }

115: PetscErrorCode ISLocate_Stride(IS is, PetscInt key, PetscInt *location)
116: {
117:   IS_Stride *sub = (IS_Stride *)is->data;
118:   PetscInt   rem, step;

120:   *location = -1;
121:   step      = sub->step;
122:   key -= sub->first;
123:   rem = key / step;
124:   if ((rem < is->map->n) && !(key % step)) *location = rem;
125:   return 0;
126: }

128: /*
129:      Returns a legitimate index memory even if
130:    the stride index set is empty.
131: */
132: PetscErrorCode ISGetIndices_Stride(IS is, const PetscInt *idx[])
133: {
134:   IS_Stride *sub = (IS_Stride *)is->data;
135:   PetscInt   i, **dx = (PetscInt **)idx;

137:   PetscMalloc1(is->map->n, (PetscInt **)idx);
138:   if (is->map->n) {
139:     (*dx)[0] = sub->first;
140:     for (i = 1; i < is->map->n; i++) (*dx)[i] = (*dx)[i - 1] + sub->step;
141:   }
142:   return 0;
143: }

145: PetscErrorCode ISRestoreIndices_Stride(IS in, const PetscInt *idx[])
146: {
147:   PetscFree(*(void **)idx);
148:   return 0;
149: }

151: PetscErrorCode ISView_Stride(IS is, PetscViewer viewer)
152: {
153:   IS_Stride        *sub = (IS_Stride *)is->data;
154:   PetscInt          i, n = is->map->n;
155:   PetscMPIInt       rank, size;
156:   PetscBool         iascii, ibinary;
157:   PetscViewerFormat fmt;

159:   PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii);
160:   PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &ibinary);
161:   if (iascii) {
162:     PetscBool matl, isperm;

164:     MPI_Comm_rank(PetscObjectComm((PetscObject)is), &rank);
165:     MPI_Comm_size(PetscObjectComm((PetscObject)is), &size);
166:     PetscViewerGetFormat(viewer, &fmt);
167:     matl = (PetscBool)(fmt == PETSC_VIEWER_ASCII_MATLAB);
168:     ISGetInfo(is, IS_PERMUTATION, IS_GLOBAL, PETSC_FALSE, &isperm);
169:     if (isperm && !matl) PetscViewerASCIIPrintf(viewer, "Index set is permutation\n");
170:     if (size == 1) {
171:       if (matl) {
172:         const char *name;

174:         PetscObjectGetName((PetscObject)is, &name);
175:         PetscViewerASCIIPrintf(viewer, "%s = [%" PetscInt_FMT " : %" PetscInt_FMT " : %" PetscInt_FMT "];\n", name, sub->first + 1, sub->step, sub->first + sub->step * (n - 1) + 1);
176:       } else {
177:         PetscViewerASCIIPrintf(viewer, "Number of indices in (stride) set %" PetscInt_FMT "\n", n);
178:         for (i = 0; i < n; i++) PetscViewerASCIIPrintf(viewer, "%" PetscInt_FMT " %" PetscInt_FMT "\n", i, sub->first + i * sub->step);
179:       }
180:       PetscViewerFlush(viewer);
181:     } else {
182:       PetscViewerASCIIPushSynchronized(viewer);
183:       if (matl) {
184:         const char *name;

186:         PetscObjectGetName((PetscObject)is, &name);
187:         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);
188:       } else {
189:         PetscViewerASCIISynchronizedPrintf(viewer, "[%d] Number of indices in (stride) set %" PetscInt_FMT "\n", rank, n);
190:         for (i = 0; i < n; i++) PetscViewerASCIISynchronizedPrintf(viewer, "[%d] %" PetscInt_FMT " %" PetscInt_FMT "\n", rank, i, sub->first + i * sub->step);
191:       }
192:       PetscViewerFlush(viewer);
193:       PetscViewerASCIIPopSynchronized(viewer);
194:     }
195:   } else if (ibinary) ISView_Binary(is, viewer);
196:   return 0;
197: }

199: PetscErrorCode ISSort_Stride(IS is)
200: {
201:   IS_Stride *sub = (IS_Stride *)is->data;

203:   if (sub->step >= 0) return 0;
204:   sub->first += (is->map->n - 1) * sub->step;
205:   sub->step *= -1;
206:   return 0;
207: }

209: PetscErrorCode ISSorted_Stride(IS is, PetscBool *flg)
210: {
211:   IS_Stride *sub = (IS_Stride *)is->data;

213:   if (sub->step >= 0) *flg = PETSC_TRUE;
214:   else *flg = PETSC_FALSE;
215:   return 0;
216: }

218: static PetscErrorCode ISUniqueLocal_Stride(IS is, PetscBool *flg)
219: {
220:   IS_Stride *sub = (IS_Stride *)is->data;

222:   if (!(is->map->n) || sub->step != 0) *flg = PETSC_TRUE;
223:   else *flg = PETSC_FALSE;
224:   return 0;
225: }

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

231:   if (!(is->map->n) || (PetscAbsInt(sub->step) == 1 && is->min == 0)) *flg = PETSC_TRUE;
232:   else *flg = PETSC_FALSE;
233:   return 0;
234: }

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

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

245: static PetscErrorCode ISOnComm_Stride(IS is, MPI_Comm comm, PetscCopyMode mode, IS *newis)
246: {
247:   IS_Stride *sub = (IS_Stride *)is->data;

249:   ISCreateStride(comm, is->map->n, sub->first, sub->step, newis);
250:   return 0;
251: }

253: static PetscErrorCode ISSetBlockSize_Stride(IS is, PetscInt bs)
254: {
255:   IS_Stride *sub = (IS_Stride *)is->data;

258:   PetscLayoutSetBlockSize(is->map, bs);
259:   return 0;
260: }

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

266:   if (sub->step == 1 && sub->first >= gstart && sub->first + is->map->n <= gend) {
267:     *start  = sub->first - gstart;
268:     *contig = PETSC_TRUE;
269:   } else {
270:     *start  = -1;
271:     *contig = PETSC_FALSE;
272:   }
273:   return 0;
274: }

276: static struct _ISOps myops = {PetscDesignatedInitializer(getindices, ISGetIndices_Stride), PetscDesignatedInitializer(restoreindices, ISRestoreIndices_Stride), PetscDesignatedInitializer(invertpermutation, ISInvertPermutation_Stride), PetscDesignatedInitializer(sort, ISSort_Stride), PetscDesignatedInitializer(sortremovedups, ISSort_Stride), PetscDesignatedInitializer(sorted, ISSorted_Stride), PetscDesignatedInitializer(duplicate, ISDuplicate_Stride), PetscDesignatedInitializer(destroy, ISDestroy_Stride), PetscDesignatedInitializer(view, ISView_Stride), PetscDesignatedInitializer(load, ISLoad_Default), PetscDesignatedInitializer(copy, ISCopy_Stride), PetscDesignatedInitializer(togeneral, ISToGeneral_Stride), PetscDesignatedInitializer(oncomm, ISOnComm_Stride), PetscDesignatedInitializer(setblocksize, ISSetBlockSize_Stride), PetscDesignatedInitializer(contiguous, ISContiguousLocal_Stride), PetscDesignatedInitializer(locate, ISLocate_Stride), PetscDesignatedInitializer(sortedlocal, ISSorted_Stride), NULL, PetscDesignatedInitializer(uniquelocal, ISUniqueLocal_Stride), NULL, PetscDesignatedInitializer(permlocal, ISPermutationLocal_Stride), NULL, PetscDesignatedInitializer(intervallocal, ISIntervalLocal_Stride), NULL};

278: /*@
279:    ISStrideSetStride - Sets the stride information for a stride index set.

281:    Logically Collective

283:    Input Parameters:
284: +  is - the index set
285: .  n - the length of the locally owned portion of the index set
286: .  first - the first element of the locally owned portion of the index set
287: -  step - the change to the next index

289:    Level: beginner

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

294: .seealso: [](sec_scatter), `IS`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`, `ISCreateStride()`, `ISStrideGetInfo()`
295: @*/
296: PetscErrorCode ISStrideSetStride(IS is, PetscInt n, PetscInt first, PetscInt step)
297: {
299:   ISClearInfoCache(is, PETSC_FALSE);
300:   PetscUseMethod(is, "ISStrideSetStride_C", (IS, PetscInt, PetscInt, PetscInt), (is, n, first, step));
301:   return 0;
302: }

304: PetscErrorCode ISStrideSetStride_Stride(IS is, PetscInt n, PetscInt first, PetscInt step)
305: {
306:   PetscInt    min, max;
307:   IS_Stride  *sub = (IS_Stride *)is->data;
308:   PetscLayout map;

310:   PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)is), n, is->map->N, is->map->bs, &map);
311:   PetscLayoutDestroy(&is->map);
312:   is->map = map;

314:   sub->first = first;
315:   sub->step  = step;
316:   if (step > 0) {
317:     min = first;
318:     max = first + step * (n - 1);
319:   } else {
320:     max = first;
321:     min = first + step * (n - 1);
322:   }

324:   is->min  = n > 0 ? min : PETSC_MAX_INT;
325:   is->max  = n > 0 ? max : PETSC_MIN_INT;
326:   is->data = (void *)sub;
327:   return 0;
328: }

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

333:    Collective

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

341:    Output Parameter:
342: .  is - the new index set

344:    Level: beginner

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

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

353: .seealso: [](sec_scatter), `IS`, `ISStrideSetStride()`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`
354: @*/
355: PetscErrorCode ISCreateStride(MPI_Comm comm, PetscInt n, PetscInt first, PetscInt step, IS *is)
356: {
357:   ISCreate(comm, is);
358:   ISSetType(*is, ISSTRIDE);
359:   ISStrideSetStride(*is, n, first, step);
360:   return 0;
361: }

363: PETSC_EXTERN PetscErrorCode ISCreate_Stride(IS is)
364: {
365:   IS_Stride *sub;

367:   PetscNew(&sub);
368:   is->data = (void *)sub;
369:   PetscMemcpy(is->ops, &myops, sizeof(myops));
370:   PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", ISStrideSetStride_Stride);
371:   PetscObjectComposeFunction((PetscObject)is, "ISShift_C", ISShift_Stride);
372:   return 0;
373: }