xref: /petsc/src/vec/is/is/impls/stride/stride.c (revision 8112c1cbf372cb53bf7c5aca994a84a6a303db4d)
1 /*
2        Index sets of evenly space integers, defined by a
3     start, stride and length.
4 */
5 #include <petsc/private/isimpl.h> /*I   "petscis.h"   I*/
6 #include <petscviewer.h>
7 
8 typedef struct {
9   PetscInt first, step;
10 } IS_Stride;
11 
ISCopy_Stride(IS is,IS isy)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;
15 
16   PetscFunctionBegin;
17   PetscCall(PetscMemcpy(isy_stride, is_stride, sizeof(IS_Stride)));
18   PetscFunctionReturn(PETSC_SUCCESS);
19 }
20 
ISShift_Stride(IS is,PetscInt shift,IS isy)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;
24 
25   PetscFunctionBegin;
26   isy_stride->first = is_stride->first + shift;
27   isy_stride->step  = is_stride->step;
28   PetscFunctionReturn(PETSC_SUCCESS);
29 }
30 
ISDuplicate_Stride(IS is,IS * newIS)31 static PetscErrorCode ISDuplicate_Stride(IS is, IS *newIS)
32 {
33   IS_Stride *sub = (IS_Stride *)is->data;
34 
35   PetscFunctionBegin;
36   PetscCall(ISCreateStride(PetscObjectComm((PetscObject)is), is->map->n, sub->first, sub->step, newIS));
37   PetscFunctionReturn(PETSC_SUCCESS);
38 }
39 
ISInvertPermutation_Stride(IS is,PetscInt nlocal,IS * perm)40 static PetscErrorCode ISInvertPermutation_Stride(IS is, PetscInt nlocal, IS *perm)
41 {
42   PetscBool isident, samelocal = (PetscBool)(nlocal == PETSC_DECIDE);
43 
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;
49 
50     if (!samelocal) {
51       n     = nlocal;
52       start = 0;
53 
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;
60 
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 }
70 
71 /*@
72   ISStrideGetInfo - Returns the first index in a stride index set and the stride width from an `IS` of `ISType` `ISSTRIDE`
73 
74   Not Collective
75 
76   Input Parameter:
77 . is - the index set
78 
79   Output Parameters:
80 + first - the first index
81 - step  - the stride width
82 
83   Level: intermediate
84 
85 .seealso: [](sec_scatter), `IS`, `ISCreateStride()`, `ISGetSize()`, `ISSTRIDE`
86 @*/
ISStrideGetInfo(IS is,PetscInt * first,PetscInt * step)87 PetscErrorCode ISStrideGetInfo(IS is, PetscInt *first, PetscInt *step)
88 {
89   IS_Stride *sub;
90   PetscBool  flg;
91 
92   PetscFunctionBegin;
93   PetscValidHeaderSpecific(is, IS_CLASSID, 1);
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");
98 
99   sub = (IS_Stride *)is->data;
100   if (first) *first = sub->first;
101   if (step) *step = sub->step;
102   PetscFunctionReturn(PETSC_SUCCESS);
103 }
104 
ISDestroy_Stride(IS is)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 }
113 
ISToGeneral_Stride(IS inis)114 static PetscErrorCode ISToGeneral_Stride(IS inis)
115 {
116   const PetscInt *idx;
117   PetscInt        n;
118 
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 }
126 
ISLocate_Stride(IS is,PetscInt key,PetscInt * location)127 static PetscErrorCode ISLocate_Stride(IS is, PetscInt key, PetscInt *location)
128 {
129   IS_Stride *sub = (IS_Stride *)is->data;
130   PetscInt   rem, step;
131 
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 }
140 
141 /*
142      Returns a legitimate index memory even if
143    the stride index set is empty.
144 */
ISGetIndices_Stride(IS is,const PetscInt * idx[])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;
149 
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 }
158 
ISRestoreIndices_Stride(IS in,const PetscInt * idx[])159 static PetscErrorCode ISRestoreIndices_Stride(IS in, const PetscInt *idx[])
160 {
161   PetscFunctionBegin;
162   PetscCall(PetscFree(*(void **)idx));
163   PetscFunctionReturn(PETSC_SUCCESS);
164 }
165 
ISView_Stride(IS is,PetscViewer viewer)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;
173 
174   PetscFunctionBegin;
175   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
176   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &ibinary));
177   if (isascii) {
178     PetscBool matl, isperm;
179 
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;
189 
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;
201 
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 }
214 
ISSort_Stride(IS is)215 static PetscErrorCode ISSort_Stride(IS is)
216 {
217   IS_Stride *sub = (IS_Stride *)is->data;
218 
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 }
225 
ISSorted_Stride(IS is,PetscBool * flg)226 static PetscErrorCode ISSorted_Stride(IS is, PetscBool *flg)
227 {
228   IS_Stride *sub = (IS_Stride *)is->data;
229 
230   PetscFunctionBegin;
231   if (sub->step >= 0) *flg = PETSC_TRUE;
232   else *flg = PETSC_FALSE;
233   PetscFunctionReturn(PETSC_SUCCESS);
234 }
235 
ISUniqueLocal_Stride(IS is,PetscBool * flg)236 static PetscErrorCode ISUniqueLocal_Stride(IS is, PetscBool *flg)
237 {
238   IS_Stride *sub = (IS_Stride *)is->data;
239 
240   PetscFunctionBegin;
241   if (!is->map->n || sub->step != 0) *flg = PETSC_TRUE;
242   else *flg = PETSC_FALSE;
243   PetscFunctionReturn(PETSC_SUCCESS);
244 }
245 
ISPermutationLocal_Stride(IS is,PetscBool * flg)246 static PetscErrorCode ISPermutationLocal_Stride(IS is, PetscBool *flg)
247 {
248   IS_Stride *sub = (IS_Stride *)is->data;
249 
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 }
255 
ISIntervalLocal_Stride(IS is,PetscBool * flg)256 static PetscErrorCode ISIntervalLocal_Stride(IS is, PetscBool *flg)
257 {
258   IS_Stride *sub = (IS_Stride *)is->data;
259 
260   PetscFunctionBegin;
261   if (!is->map->n || sub->step == 1) *flg = PETSC_TRUE;
262   else *flg = PETSC_FALSE;
263   PetscFunctionReturn(PETSC_SUCCESS);
264 }
265 
ISOnComm_Stride(IS is,MPI_Comm comm,PetscCopyMode mode,IS * newis)266 static PetscErrorCode ISOnComm_Stride(IS is, MPI_Comm comm, PetscCopyMode mode, IS *newis)
267 {
268   IS_Stride *sub = (IS_Stride *)is->data;
269 
270   PetscFunctionBegin;
271   PetscCall(ISCreateStride(comm, is->map->n, sub->first, sub->step, newis));
272   PetscFunctionReturn(PETSC_SUCCESS);
273 }
274 
ISSetBlockSize_Stride(IS is,PetscInt bs)275 static PetscErrorCode ISSetBlockSize_Stride(IS is, PetscInt bs)
276 {
277   IS_Stride *sub = (IS_Stride *)is->data;
278 
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 }
284 
ISContiguousLocal_Stride(IS is,PetscInt gstart,PetscInt gend,PetscInt * start,PetscBool * contig)285 static PetscErrorCode ISContiguousLocal_Stride(IS is, PetscInt gstart, PetscInt gend, PetscInt *start, PetscBool *contig)
286 {
287   IS_Stride *sub = (IS_Stride *)is->data;
288 
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 }
299 
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
328 
329 /*@
330   ISStrideSetStride - Sets the stride information for a stride index set.
331 
332   Logically Collective
333 
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
339 
340   Level: beginner
341 
342   Note:
343   `ISCreateStride()` can be used to create an `ISSTRIDE` and set its stride in one function call
344 
345 .seealso: [](sec_scatter), `IS`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`, `ISCreateStride()`, `ISStrideGetInfo()`
346 @*/
ISStrideSetStride(IS is,PetscInt n,PetscInt first,PetscInt step)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 }
355 
ISStrideSetStride_Stride(IS is,PetscInt n,PetscInt first,PetscInt step)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;
361 
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;
366 
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   }
376 
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 }
382 
383 /*@
384   ISCreateStride - Creates a data structure for an index set containing a list of evenly spaced integers.
385 
386   Collective
387 
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
393 
394   Output Parameter:
395 . is - the new index set
396 
397   Level: beginner
398 
399   Notes:
400   `ISStrideSetStride()` may be used to set the stride of an `ISSTRIDE` that already exists
401 
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.
405 
406 .seealso: [](sec_scatter), `IS`, `ISStrideSetStride()`, `ISCreateGeneral()`, `ISCreateBlock()`, `ISAllGather()`, `ISSTRIDE`
407 @*/
ISCreateStride(MPI_Comm comm,PetscInt n,PetscInt first,PetscInt step,IS * is)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 }
416 
ISCreate_Stride(IS is)417 PETSC_INTERN PetscErrorCode ISCreate_Stride(IS is)
418 {
419   IS_Stride *sub;
420 
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 }
429