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, MPIU_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 iascii, ibinary;
172: PetscViewerFormat fmt;
174: PetscFunctionBegin;
175: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
176: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &ibinary));
177: if (iascii) {
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: }