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:   PetscFunctionBegin;
 18:   PetscCall(PetscMemcpy(isy_stride, is_stride, sizeof(IS_Stride)));
 19:   PetscFunctionReturn(PETSC_SUCCESS);
 20: }

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

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

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

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

 41: PetscErrorCode ISInvertPermutation_Stride(IS is, PetscInt nlocal, IS *perm)
 42: {
 43:   PetscBool isident;

 45:   PetscFunctionBegin;
 46:   PetscCall(ISGetInfo(is, IS_IDENTITY, IS_GLOBAL, PETSC_TRUE, &isident));
 47:   if (isident) {
 48:     PetscInt rStart, rEnd;

 50:     PetscCall(PetscLayoutGetRange(is->map, &rStart, &rEnd));
 51:     PetscCall(ISCreateStride(PETSC_COMM_SELF, PetscMax(rEnd - rStart, 0), rStart, 1, perm));
 52:   } else {
 53:     IS              tmp;
 54:     const PetscInt *indices, n = is->map->n;

 56:     PetscCall(ISGetIndices(is, &indices));
 57:     PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is), n, indices, PETSC_COPY_VALUES, &tmp));
 58:     PetscCall(ISSetPermutation(tmp));
 59:     PetscCall(ISRestoreIndices(is, &indices));
 60:     PetscCall(ISInvertPermutation(tmp, nlocal, perm));
 61:     PetscCall(ISDestroy(&tmp));
 62:   }
 63:   PetscFunctionReturn(PETSC_SUCCESS);
 64: }

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

 69:    Not Collective

 71:    Input Parameter:
 72: .  is - the index set

 74:    Output Parameters:
 75: +  first - the first index
 76: -  step - the stride width

 78:    Level: intermediate

 80: .seealso: [](sec_scatter), `IS`, `ISCreateStride()`, `ISGetSize()`, `ISSTRIDE`
 81: @*/
 82: PetscErrorCode ISStrideGetInfo(IS is, PetscInt *first, PetscInt *step)
 83: {
 84:   IS_Stride *sub;
 85:   PetscBool  flg;

 87:   PetscFunctionBegin;
 91:   PetscCall(PetscObjectTypeCompare((PetscObject)is, ISSTRIDE, &flg));
 92:   PetscCheck(flg, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_WRONG, "IS must be of type ISSTRIDE");

 94:   sub = (IS_Stride *)is->data;
 95:   if (first) *first = sub->first;
 96:   if (step) *step = sub->step;
 97:   PetscFunctionReturn(PETSC_SUCCESS);
 98: }

100: PetscErrorCode ISDestroy_Stride(IS is)
101: {
102:   PetscFunctionBegin;
103:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", NULL));
104:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISShift_C", NULL));
105:   PetscCall(PetscFree(is->data));
106:   PetscFunctionReturn(PETSC_SUCCESS);
107: }

109: PetscErrorCode ISToGeneral_Stride(IS inis)
110: {
111:   const PetscInt *idx;
112:   PetscInt        n;

114:   PetscFunctionBegin;
115:   PetscCall(ISGetLocalSize(inis, &n));
116:   PetscCall(ISGetIndices(inis, &idx));
117:   PetscCall(ISSetType(inis, ISGENERAL));
118:   PetscCall(ISGeneralSetIndices(inis, n, idx, PETSC_OWN_POINTER));
119:   PetscFunctionReturn(PETSC_SUCCESS);
120: }

122: PetscErrorCode ISLocate_Stride(IS is, PetscInt key, PetscInt *location)
123: {
124:   IS_Stride *sub = (IS_Stride *)is->data;
125:   PetscInt   rem, step;

127:   PetscFunctionBegin;
128:   *location = -1;
129:   step      = sub->step;
130:   key -= sub->first;
131:   rem = key / step;
132:   if ((rem < is->map->n) && !(key % step)) *location = rem;
133:   PetscFunctionReturn(PETSC_SUCCESS);
134: }

136: /*
137:      Returns a legitimate index memory even if
138:    the stride index set is empty.
139: */
140: PetscErrorCode ISGetIndices_Stride(IS is, const PetscInt *idx[])
141: {
142:   IS_Stride *sub = (IS_Stride *)is->data;
143:   PetscInt   i, **dx = (PetscInt **)idx;

145:   PetscFunctionBegin;
146:   PetscCall(PetscMalloc1(is->map->n, (PetscInt **)idx));
147:   if (is->map->n) {
148:     (*dx)[0] = sub->first;
149:     for (i = 1; i < is->map->n; i++) (*dx)[i] = (*dx)[i - 1] + sub->step;
150:   }
151:   PetscFunctionReturn(PETSC_SUCCESS);
152: }

154: PetscErrorCode ISRestoreIndices_Stride(IS in, const PetscInt *idx[])
155: {
156:   PetscFunctionBegin;
157:   PetscCall(PetscFree(*(void **)idx));
158:   PetscFunctionReturn(PETSC_SUCCESS);
159: }

161: PetscErrorCode ISView_Stride(IS is, PetscViewer viewer)
162: {
163:   IS_Stride        *sub = (IS_Stride *)is->data;
164:   PetscInt          i, n = is->map->n;
165:   PetscMPIInt       rank, size;
166:   PetscBool         iascii, ibinary;
167:   PetscViewerFormat fmt;

169:   PetscFunctionBegin;
170:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
171:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &ibinary));
172:   if (iascii) {
173:     PetscBool matl, isperm;

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

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

197:         PetscCall(PetscObjectGetName((PetscObject)is, &name));
198:         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));
199:       } else {
200:         PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] Number of indices in (stride) set %" PetscInt_FMT "\n", rank, n));
201:         for (i = 0; i < n; i++) PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] %" PetscInt_FMT " %" PetscInt_FMT "\n", rank, i, sub->first + i * sub->step));
202:       }
203:       PetscCall(PetscViewerFlush(viewer));
204:       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
205:     }
206:   } else if (ibinary) PetscCall(ISView_Binary(is, viewer));
207:   PetscFunctionReturn(PETSC_SUCCESS);
208: }

210: PetscErrorCode ISSort_Stride(IS is)
211: {
212:   IS_Stride *sub = (IS_Stride *)is->data;

214:   PetscFunctionBegin;
215:   if (sub->step >= 0) PetscFunctionReturn(PETSC_SUCCESS);
216:   sub->first += (is->map->n - 1) * sub->step;
217:   sub->step *= -1;
218:   PetscFunctionReturn(PETSC_SUCCESS);
219: }

221: PetscErrorCode ISSorted_Stride(IS is, PetscBool *flg)
222: {
223:   IS_Stride *sub = (IS_Stride *)is->data;

225:   PetscFunctionBegin;
226:   if (sub->step >= 0) *flg = PETSC_TRUE;
227:   else *flg = PETSC_FALSE;
228:   PetscFunctionReturn(PETSC_SUCCESS);
229: }

231: static PetscErrorCode ISUniqueLocal_Stride(IS is, PetscBool *flg)
232: {
233:   IS_Stride *sub = (IS_Stride *)is->data;

235:   PetscFunctionBegin;
236:   if (!(is->map->n) || sub->step != 0) *flg = PETSC_TRUE;
237:   else *flg = PETSC_FALSE;
238:   PetscFunctionReturn(PETSC_SUCCESS);
239: }

241: static PetscErrorCode ISPermutationLocal_Stride(IS is, PetscBool *flg)
242: {
243:   IS_Stride *sub = (IS_Stride *)is->data;

245:   PetscFunctionBegin;
246:   if (!(is->map->n) || (PetscAbsInt(sub->step) == 1 && is->min == 0)) *flg = PETSC_TRUE;
247:   else *flg = PETSC_FALSE;
248:   PetscFunctionReturn(PETSC_SUCCESS);
249: }

251: static PetscErrorCode ISIntervalLocal_Stride(IS is, PetscBool *flg)
252: {
253:   IS_Stride *sub = (IS_Stride *)is->data;

255:   PetscFunctionBegin;
256:   if (!(is->map->n) || sub->step == 1) *flg = PETSC_TRUE;
257:   else *flg = PETSC_FALSE;
258:   PetscFunctionReturn(PETSC_SUCCESS);
259: }

261: static PetscErrorCode ISOnComm_Stride(IS is, MPI_Comm comm, PetscCopyMode mode, IS *newis)
262: {
263:   IS_Stride *sub = (IS_Stride *)is->data;

265:   PetscFunctionBegin;
266:   PetscCall(ISCreateStride(comm, is->map->n, sub->first, sub->step, newis));
267:   PetscFunctionReturn(PETSC_SUCCESS);
268: }

270: static PetscErrorCode ISSetBlockSize_Stride(IS is, PetscInt bs)
271: {
272:   IS_Stride *sub = (IS_Stride *)is->data;

274:   PetscFunctionBegin;
275:   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);
276:   PetscCall(PetscLayoutSetBlockSize(is->map, bs));
277:   PetscFunctionReturn(PETSC_SUCCESS);
278: }

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

284:   PetscFunctionBegin;
285:   if (sub->step == 1 && sub->first >= gstart && sub->first + is->map->n <= gend) {
286:     *start  = sub->first - gstart;
287:     *contig = PETSC_TRUE;
288:   } else {
289:     *start  = -1;
290:     *contig = PETSC_FALSE;
291:   }
292:   PetscFunctionReturn(PETSC_SUCCESS);
293: }

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

324: /*@
325:    ISStrideSetStride - Sets the stride information for a stride index set.

327:    Logically Collective

329:    Input Parameters:
330: +  is - the index set
331: .  n - the length of the locally owned portion of the index set
332: .  first - the first element of the locally owned portion of the index set
333: -  step - the change to the next index

335:    Level: beginner

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

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

351: PetscErrorCode ISStrideSetStride_Stride(IS is, PetscInt n, PetscInt first, PetscInt step)
352: {
353:   PetscInt    min, max;
354:   IS_Stride  *sub = (IS_Stride *)is->data;
355:   PetscLayout map;

357:   PetscFunctionBegin;
358:   PetscCall(PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)is), n, is->map->N, is->map->bs, &map));
359:   PetscCall(PetscLayoutDestroy(&is->map));
360:   is->map = map;

362:   sub->first = first;
363:   sub->step  = step;
364:   if (step > 0) {
365:     min = first;
366:     max = first + step * (n - 1);
367:   } else {
368:     max = first;
369:     min = first + step * (n - 1);
370:   }

372:   is->min  = n > 0 ? min : PETSC_MAX_INT;
373:   is->max  = n > 0 ? max : PETSC_MIN_INT;
374:   is->data = (void *)sub;
375:   PetscFunctionReturn(PETSC_SUCCESS);
376: }

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

381:    Collective

383:    Input Parameters:
384: +  comm - the MPI communicator
385: .  n - the length of the locally owned portion of the index set
386: .  first - the first element of the locally owned portion of the index set
387: -  step - the change to the next index

389:    Output Parameter:
390: .  is - the new index set

392:    Level: beginner

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

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

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

412: PETSC_EXTERN PetscErrorCode ISCreate_Stride(IS is)
413: {
414:   IS_Stride *sub;

416:   PetscFunctionBegin;
417:   PetscCall(PetscNew(&sub));
418:   is->data = (void *)sub;
419:   PetscCall(PetscMemcpy(is->ops, &myops, sizeof(myops)));
420:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", ISStrideSetStride_Stride));
421:   PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISShift_C", ISShift_Stride));
422:   PetscFunctionReturn(PETSC_SUCCESS);
423: }