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;
44: PetscFunctionBegin;
45: PetscCall(ISGetInfo(is, IS_IDENTITY, IS_GLOBAL, PETSC_TRUE, &isident));
46: if (isident) {
47: PetscInt rStart, rEnd;
49: PetscCall(PetscLayoutGetRange(is->map, &rStart, &rEnd));
50: PetscCall(ISCreateStride(PETSC_COMM_SELF, PetscMax(rEnd - rStart, 0), rStart, 1, perm));
51: } else {
52: IS tmp;
53: const PetscInt *indices, n = is->map->n;
55: PetscCall(ISGetIndices(is, &indices));
56: PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is), n, indices, PETSC_COPY_VALUES, &tmp));
57: PetscCall(ISSetPermutation(tmp));
58: PetscCall(ISRestoreIndices(is, &indices));
59: PetscCall(ISInvertPermutation(tmp, nlocal, perm));
60: PetscCall(ISDestroy(&tmp));
61: }
62: PetscFunctionReturn(PETSC_SUCCESS);
63: }
65: /*@
66: ISStrideGetInfo - Returns the first index in a stride index set and the stride width from an `IS` of `ISType` `ISSTRIDE`
68: Not Collective
70: Input Parameter:
71: . is - the index set
73: Output Parameters:
74: + first - the first index
75: - step - the stride width
77: Level: intermediate
79: .seealso: [](sec_scatter), `IS`, `ISCreateStride()`, `ISGetSize()`, `ISSTRIDE`
80: @*/
81: PetscErrorCode ISStrideGetInfo(IS is, PetscInt *first, PetscInt *step)
82: {
83: IS_Stride *sub;
84: PetscBool flg;
86: PetscFunctionBegin;
88: if (first) PetscAssertPointer(first, 2);
89: if (step) PetscAssertPointer(step, 3);
90: PetscCall(PetscObjectTypeCompare((PetscObject)is, ISSTRIDE, &flg));
91: PetscCheck(flg, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_WRONG, "IS must be of type ISSTRIDE");
93: sub = (IS_Stride *)is->data;
94: if (first) *first = sub->first;
95: if (step) *step = sub->step;
96: PetscFunctionReturn(PETSC_SUCCESS);
97: }
99: static PetscErrorCode ISDestroy_Stride(IS is)
100: {
101: PetscFunctionBegin;
102: PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", NULL));
103: PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISShift_C", NULL));
104: PetscCall(PetscFree(is->data));
105: PetscFunctionReturn(PETSC_SUCCESS);
106: }
108: static PetscErrorCode ISToGeneral_Stride(IS inis)
109: {
110: const PetscInt *idx;
111: PetscInt n;
113: PetscFunctionBegin;
114: PetscCall(ISGetLocalSize(inis, &n));
115: PetscCall(ISGetIndices(inis, &idx));
116: PetscCall(ISSetType(inis, ISGENERAL));
117: PetscCall(ISGeneralSetIndices(inis, n, idx, PETSC_OWN_POINTER));
118: PetscFunctionReturn(PETSC_SUCCESS);
119: }
121: static PetscErrorCode ISLocate_Stride(IS is, PetscInt key, PetscInt *location)
122: {
123: IS_Stride *sub = (IS_Stride *)is->data;
124: PetscInt rem, step;
126: PetscFunctionBegin;
127: *location = -1;
128: step = sub->step;
129: key -= sub->first;
130: rem = key / step;
131: if ((rem < is->map->n) && !(key % step)) *location = rem;
132: PetscFunctionReturn(PETSC_SUCCESS);
133: }
135: /*
136: Returns a legitimate index memory even if
137: the stride index set is empty.
138: */
139: static PetscErrorCode ISGetIndices_Stride(IS is, const PetscInt *idx[])
140: {
141: IS_Stride *sub = (IS_Stride *)is->data;
142: PetscInt i, **dx = (PetscInt **)idx;
144: PetscFunctionBegin;
145: PetscCall(PetscMalloc1(is->map->n, (PetscInt **)idx));
146: if (is->map->n) {
147: (*dx)[0] = sub->first;
148: for (i = 1; i < is->map->n; i++) (*dx)[i] = (*dx)[i - 1] + sub->step;
149: }
150: PetscFunctionReturn(PETSC_SUCCESS);
151: }
153: static PetscErrorCode ISRestoreIndices_Stride(IS in, const PetscInt *idx[])
154: {
155: PetscFunctionBegin;
156: PetscCall(PetscFree(*(void **)idx));
157: PetscFunctionReturn(PETSC_SUCCESS);
158: }
160: static PetscErrorCode ISView_Stride(IS is, PetscViewer viewer)
161: {
162: IS_Stride *sub = (IS_Stride *)is->data;
163: PetscInt i, n = is->map->n;
164: PetscMPIInt rank, size;
165: PetscBool iascii, ibinary;
166: PetscViewerFormat fmt;
168: PetscFunctionBegin;
169: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
170: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &ibinary));
171: if (iascii) {
172: PetscBool matl, isperm;
174: PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)is), &rank));
175: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)is), &size));
176: PetscCall(PetscViewerGetFormat(viewer, &fmt));
177: matl = (PetscBool)(fmt == PETSC_VIEWER_ASCII_MATLAB);
178: PetscCall(ISGetInfo(is, IS_PERMUTATION, IS_GLOBAL, PETSC_FALSE, &isperm));
179: if (isperm && !matl) PetscCall(PetscViewerASCIIPrintf(viewer, "Index set is permutation\n"));
180: if (size == 1) {
181: if (matl) {
182: const char *name;
184: PetscCall(PetscObjectGetName((PetscObject)is, &name));
185: PetscCall(PetscViewerASCIIPrintf(viewer, "%s = [%" PetscInt_FMT " : %" PetscInt_FMT " : %" PetscInt_FMT "];\n", name, sub->first + 1, sub->step, sub->first + sub->step * (n - 1) + 1));
186: } else {
187: PetscCall(PetscViewerASCIIPrintf(viewer, "Number of indices in (stride) set %" PetscInt_FMT "\n", n));
188: for (i = 0; i < n; i++) PetscCall(PetscViewerASCIIPrintf(viewer, "%" PetscInt_FMT " %" PetscInt_FMT "\n", i, sub->first + i * sub->step));
189: }
190: PetscCall(PetscViewerFlush(viewer));
191: } else {
192: PetscCall(PetscViewerASCIIPushSynchronized(viewer));
193: if (matl) {
194: const char *name;
196: PetscCall(PetscObjectGetName((PetscObject)is, &name));
197: 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));
198: } else {
199: PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] Number of indices in (stride) set %" PetscInt_FMT "\n", rank, n));
200: for (i = 0; i < n; i++) PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] %" PetscInt_FMT " %" PetscInt_FMT "\n", rank, i, sub->first + i * sub->step));
201: }
202: PetscCall(PetscViewerFlush(viewer));
203: PetscCall(PetscViewerASCIIPopSynchronized(viewer));
204: }
205: } else if (ibinary) PetscCall(ISView_Binary(is, viewer));
206: PetscFunctionReturn(PETSC_SUCCESS);
207: }
209: static PetscErrorCode ISSort_Stride(IS is)
210: {
211: IS_Stride *sub = (IS_Stride *)is->data;
213: PetscFunctionBegin;
214: if (sub->step >= 0) PetscFunctionReturn(PETSC_SUCCESS);
215: sub->first += (is->map->n - 1) * sub->step;
216: sub->step *= -1;
217: PetscFunctionReturn(PETSC_SUCCESS);
218: }
220: static PetscErrorCode ISSorted_Stride(IS is, PetscBool *flg)
221: {
222: IS_Stride *sub = (IS_Stride *)is->data;
224: PetscFunctionBegin;
225: if (sub->step >= 0) *flg = PETSC_TRUE;
226: else *flg = PETSC_FALSE;
227: PetscFunctionReturn(PETSC_SUCCESS);
228: }
230: static PetscErrorCode ISUniqueLocal_Stride(IS is, PetscBool *flg)
231: {
232: IS_Stride *sub = (IS_Stride *)is->data;
234: PetscFunctionBegin;
235: if (!is->map->n || sub->step != 0) *flg = PETSC_TRUE;
236: else *flg = PETSC_FALSE;
237: PetscFunctionReturn(PETSC_SUCCESS);
238: }
240: static PetscErrorCode ISPermutationLocal_Stride(IS is, PetscBool *flg)
241: {
242: IS_Stride *sub = (IS_Stride *)is->data;
244: PetscFunctionBegin;
245: if (!is->map->n || (PetscAbsInt(sub->step) == 1 && is->min == 0)) *flg = PETSC_TRUE;
246: else *flg = PETSC_FALSE;
247: PetscFunctionReturn(PETSC_SUCCESS);
248: }
250: static PetscErrorCode ISIntervalLocal_Stride(IS is, PetscBool *flg)
251: {
252: IS_Stride *sub = (IS_Stride *)is->data;
254: PetscFunctionBegin;
255: if (!is->map->n || sub->step == 1) *flg = PETSC_TRUE;
256: else *flg = PETSC_FALSE;
257: PetscFunctionReturn(PETSC_SUCCESS);
258: }
260: static PetscErrorCode ISOnComm_Stride(IS is, MPI_Comm comm, PetscCopyMode mode, IS *newis)
261: {
262: IS_Stride *sub = (IS_Stride *)is->data;
264: PetscFunctionBegin;
265: PetscCall(ISCreateStride(comm, is->map->n, sub->first, sub->step, newis));
266: PetscFunctionReturn(PETSC_SUCCESS);
267: }
269: static PetscErrorCode ISSetBlockSize_Stride(IS is, PetscInt bs)
270: {
271: IS_Stride *sub = (IS_Stride *)is->data;
273: PetscFunctionBegin;
274: 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);
275: PetscCall(PetscLayoutSetBlockSize(is->map, bs));
276: PetscFunctionReturn(PETSC_SUCCESS);
277: }
279: static PetscErrorCode ISContiguousLocal_Stride(IS is, PetscInt gstart, PetscInt gend, PetscInt *start, PetscBool *contig)
280: {
281: IS_Stride *sub = (IS_Stride *)is->data;
283: PetscFunctionBegin;
284: if (sub->step == 1 && sub->first >= gstart && sub->first + is->map->n <= gend) {
285: *start = sub->first - gstart;
286: *contig = PETSC_TRUE;
287: } else {
288: *start = -1;
289: *contig = PETSC_FALSE;
290: }
291: PetscFunctionReturn(PETSC_SUCCESS);
292: }
294: // clang-format off
295: static const struct _ISOps myops = {
296: PetscDesignatedInitializer(getindices, ISGetIndices_Stride),
297: PetscDesignatedInitializer(restoreindices, ISRestoreIndices_Stride),
298: PetscDesignatedInitializer(invertpermutation, ISInvertPermutation_Stride),
299: PetscDesignatedInitializer(sort, ISSort_Stride),
300: PetscDesignatedInitializer(sortremovedups, ISSort_Stride),
301: PetscDesignatedInitializer(sorted, ISSorted_Stride),
302: PetscDesignatedInitializer(duplicate, ISDuplicate_Stride),
303: PetscDesignatedInitializer(destroy, ISDestroy_Stride),
304: PetscDesignatedInitializer(view, ISView_Stride),
305: PetscDesignatedInitializer(load, ISLoad_Default),
306: PetscDesignatedInitializer(copy, ISCopy_Stride),
307: PetscDesignatedInitializer(togeneral, ISToGeneral_Stride),
308: PetscDesignatedInitializer(oncomm, ISOnComm_Stride),
309: PetscDesignatedInitializer(setblocksize, ISSetBlockSize_Stride),
310: PetscDesignatedInitializer(contiguous, ISContiguousLocal_Stride),
311: PetscDesignatedInitializer(locate, ISLocate_Stride),
312: PetscDesignatedInitializer(sortedlocal, ISSorted_Stride),
313: PetscDesignatedInitializer(sortedglobal, NULL),
314: PetscDesignatedInitializer(uniquelocal, ISUniqueLocal_Stride),
315: PetscDesignatedInitializer(uniqueglobal, NULL),
316: PetscDesignatedInitializer(permlocal, ISPermutationLocal_Stride),
317: PetscDesignatedInitializer(permglobal, NULL),
318: PetscDesignatedInitializer(intervallocal, ISIntervalLocal_Stride),
319: PetscDesignatedInitializer(intervalglobal, NULL)
320: };
321: // clang-format on
323: /*@
324: ISStrideSetStride - Sets the stride information for a stride index set.
326: Logically Collective
328: Input Parameters:
329: + is - the index set
330: . n - the length of the locally owned portion of the index set
331: . first - the first element of the locally owned portion of the index set
332: - step - the change to the next index
334: Level: beginner
336: Note:
337: `ISCreateStride()` can be used to create an `ISSTRIDE` and set its stride in one function call
339: .seealso: [](sec_scatter), `IS`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`, `ISCreateStride()`, `ISStrideGetInfo()`
340: @*/
341: PetscErrorCode ISStrideSetStride(IS is, PetscInt n, PetscInt first, PetscInt step)
342: {
343: PetscFunctionBegin;
344: PetscCheck(n >= 0, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_OUTOFRANGE, "Negative length %" PetscInt_FMT " not valid", n);
345: PetscCall(ISClearInfoCache(is, PETSC_FALSE));
346: PetscUseMethod(is, "ISStrideSetStride_C", (IS, PetscInt, PetscInt, PetscInt), (is, n, first, step));
347: PetscFunctionReturn(PETSC_SUCCESS);
348: }
350: static PetscErrorCode ISStrideSetStride_Stride(IS is, PetscInt n, PetscInt first, PetscInt step)
351: {
352: PetscInt min, max;
353: IS_Stride *sub = (IS_Stride *)is->data;
354: PetscLayout map;
356: PetscFunctionBegin;
357: PetscCall(PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)is), n, is->map->N, is->map->bs, &map));
358: PetscCall(PetscLayoutDestroy(&is->map));
359: is->map = map;
361: sub->first = first;
362: sub->step = step;
363: if (step > 0) {
364: min = first;
365: max = first + step * (n - 1);
366: } else {
367: max = first;
368: min = first + step * (n - 1);
369: }
371: is->min = n > 0 ? min : PETSC_MAX_INT;
372: is->max = n > 0 ? max : PETSC_MIN_INT;
373: is->data = (void *)sub;
374: PetscFunctionReturn(PETSC_SUCCESS);
375: }
377: /*@
378: ISCreateStride - Creates a data structure for an index set containing a list of evenly spaced integers.
380: Collective
382: Input Parameters:
383: + comm - the MPI communicator
384: . n - the length of the locally owned portion of the index set
385: . first - the first element of the locally owned portion of the index set
386: - step - the change to the next index
388: Output Parameter:
389: . is - the new index set
391: Level: beginner
393: Notes:
394: `ISStrideSetStride()` may be used to set the stride of an `ISSTRIDE` that already exists
396: When the communicator is not `MPI_COMM_SELF`, the operations on `IS` are NOT
397: conceptually the same as `MPI_Group` operations. The `IS` are the
398: distributed sets of indices and thus certain operations on them are collective.
400: .seealso: [](sec_scatter), `IS`, `ISStrideSetStride()`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`
401: @*/
402: PetscErrorCode ISCreateStride(MPI_Comm comm, PetscInt n, PetscInt first, PetscInt step, IS *is)
403: {
404: PetscFunctionBegin;
405: PetscCall(ISCreate(comm, is));
406: PetscCall(ISSetType(*is, ISSTRIDE));
407: PetscCall(ISStrideSetStride(*is, n, first, step));
408: PetscFunctionReturn(PETSC_SUCCESS);
409: }
411: PETSC_INTERN PetscErrorCode ISCreate_Stride(IS is)
412: {
413: IS_Stride *sub;
415: PetscFunctionBegin;
416: PetscCall(PetscNew(&sub));
417: is->data = (void *)sub;
418: is->ops[0] = myops;
419: PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISStrideSetStride_C", ISStrideSetStride_Stride));
420: PetscCall(PetscObjectComposeFunction((PetscObject)is, "ISShift_C", ISShift_Stride));
421: PetscFunctionReturn(PETSC_SUCCESS);
422: }