xref: /petsc/src/vec/is/ao/impls/memscalable/aomemscalable.c (revision 6497c311e7b976d467be1503c1effce92a60525c)
11447629fSBarry Smith /*
21447629fSBarry Smith     The memory scalable AO application ordering routines. These store the
3d38ec673SBarry Smith   orderings on each process for that process' range of values, this is more memory-efficient than `AOBASIC`
41447629fSBarry Smith */
51447629fSBarry Smith 
61447629fSBarry Smith #include <../src/vec/is/ao/aoimpl.h> /*I  "petscao.h"   I*/
71447629fSBarry Smith 
81447629fSBarry Smith typedef struct {
91447629fSBarry Smith   PetscInt   *app_loc;   /* app_loc[i] is the partner for the ith local PETSc slot */
101447629fSBarry Smith   PetscInt   *petsc_loc; /* petsc_loc[j] is the partner for the jth local app slot */
111447629fSBarry Smith   PetscLayout map;       /* determines the local sizes of ao */
121447629fSBarry Smith } AO_MemoryScalable;
131447629fSBarry Smith 
141447629fSBarry Smith /*
156bd6ae52SBarry Smith        All processors ship the data to process 0 to be printed; note that this is not scalable because
166bd6ae52SBarry Smith        process 0 allocates space for all the orderings entry across all the processes
171447629fSBarry Smith */
18da8c939bSJacob Faibussowitsch static PetscErrorCode AOView_MemoryScalable(AO ao, PetscViewer viewer)
19d71ae5a4SJacob Faibussowitsch {
201447629fSBarry Smith   PetscMPIInt        rank, size;
211447629fSBarry Smith   AO_MemoryScalable *aomems = (AO_MemoryScalable *)ao->data;
221447629fSBarry Smith   PetscBool          iascii;
231447629fSBarry Smith   PetscMPIInt        tag_app, tag_petsc;
241447629fSBarry Smith   PetscLayout        map = aomems->map;
25*6497c311SBarry Smith   PetscInt          *app, *app_loc, *petsc, *petsc_loc, len, j;
261447629fSBarry Smith   MPI_Status         status;
271447629fSBarry Smith 
281447629fSBarry Smith   PetscFunctionBegin;
299566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
3028b400f6SJacob Faibussowitsch   PetscCheck(iascii, PetscObjectComm((PetscObject)viewer), PETSC_ERR_SUP, "Viewer type %s not supported for AO MemoryScalable", ((PetscObject)viewer)->type_name);
311447629fSBarry Smith 
329566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)ao), &rank));
339566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)ao), &size));
341447629fSBarry Smith 
359566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)ao, &tag_app));
369566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)ao, &tag_petsc));
371447629fSBarry Smith 
38dd400576SPatrick Sanan   if (rank == 0) {
399566063dSJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(viewer, "Number of elements in ordering %" PetscInt_FMT "\n", ao->N));
409566063dSJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(viewer, "PETSc->App  App->PETSc\n"));
411447629fSBarry Smith 
429566063dSJacob Faibussowitsch     PetscCall(PetscMalloc2(map->N, &app, map->N, &petsc));
431447629fSBarry Smith     len = map->n;
441447629fSBarry Smith     /* print local AO */
459566063dSJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(viewer, "Process [%d]\n", rank));
46*6497c311SBarry Smith     for (PetscInt i = 0; i < len; i++) PetscCall(PetscViewerASCIIPrintf(viewer, "%3" PetscInt_FMT "  %3" PetscInt_FMT "    %3" PetscInt_FMT "  %3" PetscInt_FMT "\n", i, aomems->app_loc[i], i, aomems->petsc_loc[i]));
471447629fSBarry Smith 
481447629fSBarry Smith     /* recv and print off-processor's AO */
49*6497c311SBarry Smith     for (PetscMPIInt i = 1; i < size; i++) {
501447629fSBarry Smith       len       = map->range[i + 1] - map->range[i];
511447629fSBarry Smith       app_loc   = app + map->range[i];
521447629fSBarry Smith       petsc_loc = petsc + map->range[i];
53*6497c311SBarry Smith       PetscCallMPI(MPIU_Recv(app_loc, len, MPIU_INT, i, tag_app, PetscObjectComm((PetscObject)ao), &status));
54*6497c311SBarry Smith       PetscCallMPI(MPIU_Recv(petsc_loc, len, MPIU_INT, i, tag_petsc, PetscObjectComm((PetscObject)ao), &status));
55*6497c311SBarry Smith       PetscCall(PetscViewerASCIIPrintf(viewer, "Process [%d]\n", i));
5648a46eb9SPierre Jolivet       for (j = 0; j < len; j++) PetscCall(PetscViewerASCIIPrintf(viewer, "%3" PetscInt_FMT "  %3" PetscInt_FMT "    %3" PetscInt_FMT "  %3" PetscInt_FMT "\n", map->range[i] + j, app_loc[j], map->range[i] + j, petsc_loc[j]));
571447629fSBarry Smith     }
589566063dSJacob Faibussowitsch     PetscCall(PetscFree2(app, petsc));
591447629fSBarry Smith 
601447629fSBarry Smith   } else {
611447629fSBarry Smith     /* send values */
62*6497c311SBarry Smith     PetscCallMPI(MPIU_Send((void *)aomems->app_loc, map->n, MPIU_INT, 0, tag_app, PetscObjectComm((PetscObject)ao)));
63*6497c311SBarry Smith     PetscCallMPI(MPIU_Send((void *)aomems->petsc_loc, map->n, MPIU_INT, 0, tag_petsc, PetscObjectComm((PetscObject)ao)));
641447629fSBarry Smith   }
659566063dSJacob Faibussowitsch   PetscCall(PetscViewerFlush(viewer));
663ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
671447629fSBarry Smith }
681447629fSBarry Smith 
69da8c939bSJacob Faibussowitsch static PetscErrorCode AODestroy_MemoryScalable(AO ao)
70d71ae5a4SJacob Faibussowitsch {
711447629fSBarry Smith   AO_MemoryScalable *aomems = (AO_MemoryScalable *)ao->data;
721447629fSBarry Smith 
731447629fSBarry Smith   PetscFunctionBegin;
749566063dSJacob Faibussowitsch   PetscCall(PetscFree2(aomems->app_loc, aomems->petsc_loc));
759566063dSJacob Faibussowitsch   PetscCall(PetscLayoutDestroy(&aomems->map));
769566063dSJacob Faibussowitsch   PetscCall(PetscFree(aomems));
773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
781447629fSBarry Smith }
791447629fSBarry Smith 
801447629fSBarry Smith /*
811447629fSBarry Smith    Input Parameters:
821447629fSBarry Smith +   ao - the application ordering context
831447629fSBarry Smith .   n  - the number of integers in ia[]
841447629fSBarry Smith .   ia - the integers; these are replaced with their mapped value
851447629fSBarry Smith -   maploc - app_loc or petsc_loc in struct "AO_MemoryScalable"
861447629fSBarry Smith 
871447629fSBarry Smith    Output Parameter:
881447629fSBarry Smith .   ia - the mapped interges
891447629fSBarry Smith  */
90da8c939bSJacob Faibussowitsch static PetscErrorCode AOMap_MemoryScalable_private(AO ao, PetscInt n, PetscInt *ia, const PetscInt *maploc)
91d71ae5a4SJacob Faibussowitsch {
921447629fSBarry Smith   AO_MemoryScalable *aomems = (AO_MemoryScalable *)ao->data;
931447629fSBarry Smith   MPI_Comm           comm;
94*6497c311SBarry Smith   PetscMPIInt        rank, size, tag1, tag2, count;
95*6497c311SBarry Smith   PetscMPIInt       *owner, j;
96*6497c311SBarry Smith   PetscInt           nmax, *sindices, *rindices, idx, lastidx, *sindices2, *rindices2, *sizes, *start;
97*6497c311SBarry Smith   PetscMPIInt        nreceives, nsends;
986bd6ae52SBarry Smith   const PetscInt    *owners = aomems->map->range;
991447629fSBarry Smith   MPI_Request       *send_waits, *recv_waits, *send_waits2, *recv_waits2;
1001447629fSBarry Smith   MPI_Status         recv_status;
101*6497c311SBarry Smith   PetscMPIInt        source, widx;
102*6497c311SBarry Smith   PetscCount         nindices;
103*6497c311SBarry Smith   PetscInt          *rbuf, *sbuf, inreceives;
1041447629fSBarry Smith   MPI_Status        *send_status, *send_status2;
1051447629fSBarry Smith 
1061447629fSBarry Smith   PetscFunctionBegin;
1079566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)ao, &comm));
1089566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1099566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
1101447629fSBarry Smith 
1111447629fSBarry Smith   /*  first count number of contributors to each processor */
1129566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(size, &start));
1139566063dSJacob Faibussowitsch   PetscCall(PetscCalloc2(2 * size, &sizes, n, &owner));
1141447629fSBarry Smith 
1151447629fSBarry Smith   j       = 0;
1161447629fSBarry Smith   lastidx = -1;
117*6497c311SBarry Smith   for (PetscInt i = 0; i < n; i++) {
1186bd6ae52SBarry Smith     if (ia[i] < 0) owner[i] = -1;      /* mark negative entries (which are not to be mapped) with a special negative value */
1196bd6ae52SBarry Smith     if (ia[i] >= ao->N) owner[i] = -2; /* mark out of range entries with special negative value */
1206bd6ae52SBarry Smith     else {
1211447629fSBarry Smith       /* if indices are NOT locally sorted, need to start search at the beginning */
1221447629fSBarry Smith       if (lastidx > (idx = ia[i])) j = 0;
1231447629fSBarry Smith       lastidx = idx;
1241447629fSBarry Smith       for (; j < size; j++) {
1251447629fSBarry Smith         if (idx >= owners[j] && idx < owners[j + 1]) {
12676ec1555SBarry Smith           sizes[2 * j]++;       /* num of indices to be sent */
12776ec1555SBarry Smith           sizes[2 * j + 1] = 1; /* send to proc[j] */
1281447629fSBarry Smith           owner[i]         = j;
1291447629fSBarry Smith           break;
1301447629fSBarry Smith         }
1311447629fSBarry Smith       }
1321447629fSBarry Smith     }
1336bd6ae52SBarry Smith   }
13476ec1555SBarry Smith   sizes[2 * rank] = sizes[2 * rank + 1] = 0; /* do not receive from self! */
1351447629fSBarry Smith   nsends                                = 0;
136*6497c311SBarry Smith   for (PetscMPIInt i = 0; i < size; i++) nsends += sizes[2 * i + 1];
1371447629fSBarry Smith 
1381447629fSBarry Smith   /* inform other processors of number of messages and max length*/
139*6497c311SBarry Smith   PetscCall(PetscMaxSum(comm, sizes, &nmax, &inreceives));
140*6497c311SBarry Smith   PetscCall(PetscMPIIntCast(inreceives, &nreceives));
1411447629fSBarry Smith 
1421447629fSBarry Smith   /* allocate arrays */
1439566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)ao, &tag1));
1449566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)ao, &tag2));
1451447629fSBarry Smith 
1469566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(nreceives * nmax, &rindices, nreceives, &recv_waits));
1479566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(nsends * nmax, &rindices2, nsends, &recv_waits2));
1481447629fSBarry Smith 
1499566063dSJacob Faibussowitsch   PetscCall(PetscMalloc3(n, &sindices, nsends, &send_waits, nsends, &send_status));
1509566063dSJacob Faibussowitsch   PetscCall(PetscMalloc3(n, &sindices2, nreceives, &send_waits2, nreceives, &send_status2));
1511447629fSBarry Smith 
1521447629fSBarry Smith   /* post 1st receives: receive others requests
1531447629fSBarry Smith      since we don't know how long each individual message is we
1541447629fSBarry Smith      allocate the largest needed buffer for each receive. Potentially
1551447629fSBarry Smith      this is a lot of wasted space.
1561447629fSBarry Smith   */
157*6497c311SBarry Smith   for (PetscMPIInt i = 0; i < nreceives; i++) PetscCallMPI(MPIU_Irecv(rindices + nmax * i, nmax, MPIU_INT, MPI_ANY_SOURCE, tag1, comm, recv_waits + i));
1581447629fSBarry Smith 
1591447629fSBarry Smith   /* do 1st sends:
1601447629fSBarry Smith       1) starts[i] gives the starting index in svalues for stuff going to
1611447629fSBarry Smith          the ith processor
1621447629fSBarry Smith   */
1631447629fSBarry Smith   start[0] = 0;
164*6497c311SBarry Smith   for (PetscMPIInt i = 1; i < size; i++) start[i] = start[i - 1] + sizes[2 * i - 2];
165*6497c311SBarry Smith   for (PetscInt i = 0; i < n; i++) {
1661447629fSBarry Smith     j = owner[i];
1676bd6ae52SBarry Smith     if (j == -1) continue; /* do not remap negative entries in ia[] */
1689371c9d4SSatish Balay     else if (j == -2) { /* out of range entries get mapped to -1 */ ia[i] = -1;
1696bd6ae52SBarry Smith       continue;
1706bd6ae52SBarry Smith     } else if (j != rank) {
1711447629fSBarry Smith       sindices[start[j]++] = ia[i];
1721447629fSBarry Smith     } else { /* compute my own map */
1731447629fSBarry Smith       ia[i] = maploc[ia[i] - owners[rank]];
1741447629fSBarry Smith     }
1751447629fSBarry Smith   }
1761447629fSBarry Smith 
1771447629fSBarry Smith   start[0] = 0;
178*6497c311SBarry Smith   for (PetscMPIInt i = 1; i < size; i++) start[i] = start[i - 1] + sizes[2 * i - 2];
179*6497c311SBarry Smith   count = 0;
180*6497c311SBarry Smith   for (PetscMPIInt i = 0; i < size; i++) {
18176ec1555SBarry Smith     if (sizes[2 * i + 1]) {
1821447629fSBarry Smith       /* send my request to others */
183*6497c311SBarry Smith       PetscCallMPI(MPIU_Isend(sindices + start[i], sizes[2 * i], MPIU_INT, i, tag1, comm, send_waits + count));
1841447629fSBarry Smith       /* post receive for the answer of my request */
185*6497c311SBarry Smith       PetscCallMPI(MPIU_Irecv(sindices2 + start[i], sizes[2 * i], MPIU_INT, i, tag2, comm, recv_waits2 + count));
1861447629fSBarry Smith       count++;
1871447629fSBarry Smith     }
1881447629fSBarry Smith   }
189*6497c311SBarry Smith   PetscCheck(nsends == count, comm, PETSC_ERR_SUP, "nsends %d != count %d", nsends, count);
1901447629fSBarry Smith 
1911447629fSBarry Smith   /* wait on 1st sends */
192*6497c311SBarry Smith   if (nsends) PetscCallMPI(MPI_Waitall((PetscMPIInt)nsends, send_waits, send_status));
1931447629fSBarry Smith 
1941447629fSBarry Smith   /* 1st recvs: other's requests */
1951447629fSBarry Smith   for (j = 0; j < nreceives; j++) {
196*6497c311SBarry Smith     PetscCallMPI(MPI_Waitany((PetscMPIInt)nreceives, recv_waits, &widx, &recv_status)); /* idx: index of handle for operation that completed */
197*6497c311SBarry Smith     PetscCallMPI(MPIU_Get_count(&recv_status, MPIU_INT, &nindices));
1981447629fSBarry Smith     rbuf   = rindices + nmax * widx; /* global index */
1991447629fSBarry Smith     source = recv_status.MPI_SOURCE;
2001447629fSBarry Smith 
2011447629fSBarry Smith     /* compute mapping */
2021447629fSBarry Smith     sbuf = rbuf;
203*6497c311SBarry Smith     for (PetscCount i = 0; i < nindices; i++) sbuf[i] = maploc[rbuf[i] - owners[rank]];
2041447629fSBarry Smith 
2051447629fSBarry Smith     /* send mapping back to the sender */
206*6497c311SBarry Smith     PetscCallMPI(MPIU_Isend(sbuf, nindices, MPIU_INT, source, tag2, comm, send_waits2 + widx));
2071447629fSBarry Smith   }
2081447629fSBarry Smith 
2091447629fSBarry Smith   /* wait on 2nd sends */
210*6497c311SBarry Smith   if (nreceives) PetscCallMPI(MPI_Waitall((PetscMPIInt)nreceives, send_waits2, send_status2));
2111447629fSBarry Smith 
2121447629fSBarry Smith   /* 2nd recvs: for the answer of my request */
2131447629fSBarry Smith   for (j = 0; j < nsends; j++) {
214*6497c311SBarry Smith     PetscCallMPI(MPI_Waitany((PetscMPIInt)nsends, recv_waits2, &widx, &recv_status));
215*6497c311SBarry Smith     PetscCallMPI(MPIU_Get_count(&recv_status, MPIU_INT, &nindices));
2161447629fSBarry Smith     source = recv_status.MPI_SOURCE;
2171447629fSBarry Smith     /* pack output ia[] */
2181447629fSBarry Smith     rbuf  = sindices2 + start[source];
2191447629fSBarry Smith     count = 0;
220*6497c311SBarry Smith     for (PetscCount i = 0; i < n; i++) {
2211447629fSBarry Smith       if (source == owner[i]) ia[i] = rbuf[count++];
2221447629fSBarry Smith     }
2231447629fSBarry Smith   }
2241447629fSBarry Smith 
2251447629fSBarry Smith   /* free arrays */
2269566063dSJacob Faibussowitsch   PetscCall(PetscFree(start));
2279566063dSJacob Faibussowitsch   PetscCall(PetscFree2(sizes, owner));
2289566063dSJacob Faibussowitsch   PetscCall(PetscFree2(rindices, recv_waits));
2299566063dSJacob Faibussowitsch   PetscCall(PetscFree2(rindices2, recv_waits2));
2309566063dSJacob Faibussowitsch   PetscCall(PetscFree3(sindices, send_waits, send_status));
2319566063dSJacob Faibussowitsch   PetscCall(PetscFree3(sindices2, send_waits2, send_status2));
2323ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2331447629fSBarry Smith }
2341447629fSBarry Smith 
235da8c939bSJacob Faibussowitsch static PetscErrorCode AOPetscToApplication_MemoryScalable(AO ao, PetscInt n, PetscInt *ia)
236d71ae5a4SJacob Faibussowitsch {
2371447629fSBarry Smith   AO_MemoryScalable *aomems  = (AO_MemoryScalable *)ao->data;
2381447629fSBarry Smith   PetscInt          *app_loc = aomems->app_loc;
2391447629fSBarry Smith 
2401447629fSBarry Smith   PetscFunctionBegin;
2419566063dSJacob Faibussowitsch   PetscCall(AOMap_MemoryScalable_private(ao, n, ia, app_loc));
2423ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2431447629fSBarry Smith }
2441447629fSBarry Smith 
245da8c939bSJacob Faibussowitsch static PetscErrorCode AOApplicationToPetsc_MemoryScalable(AO ao, PetscInt n, PetscInt *ia)
246d71ae5a4SJacob Faibussowitsch {
2471447629fSBarry Smith   AO_MemoryScalable *aomems    = (AO_MemoryScalable *)ao->data;
2481447629fSBarry Smith   PetscInt          *petsc_loc = aomems->petsc_loc;
2491447629fSBarry Smith 
2501447629fSBarry Smith   PetscFunctionBegin;
2519566063dSJacob Faibussowitsch   PetscCall(AOMap_MemoryScalable_private(ao, n, ia, petsc_loc));
2523ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2531447629fSBarry Smith }
2541447629fSBarry Smith 
255da8c939bSJacob Faibussowitsch static const struct _AOOps AOOps_MemoryScalable = {
256267267bdSJacob Faibussowitsch   PetscDesignatedInitializer(view, AOView_MemoryScalable),
257267267bdSJacob Faibussowitsch   PetscDesignatedInitializer(destroy, AODestroy_MemoryScalable),
258267267bdSJacob Faibussowitsch   PetscDesignatedInitializer(petsctoapplication, AOPetscToApplication_MemoryScalable),
259267267bdSJacob Faibussowitsch   PetscDesignatedInitializer(applicationtopetsc, AOApplicationToPetsc_MemoryScalable),
2601447629fSBarry Smith };
2611447629fSBarry Smith 
262da8c939bSJacob Faibussowitsch static PetscErrorCode AOCreateMemoryScalable_private(MPI_Comm comm, PetscInt napp, const PetscInt from_array[], const PetscInt to_array[], AO ao, PetscInt *aomap_loc)
263d71ae5a4SJacob Faibussowitsch {
2641447629fSBarry Smith   AO_MemoryScalable *aomems  = (AO_MemoryScalable *)ao->data;
2651447629fSBarry Smith   PetscLayout        map     = aomems->map;
2661447629fSBarry Smith   PetscInt           n_local = map->n, i, j;
2671447629fSBarry Smith   PetscMPIInt        rank, size, tag;
26876ec1555SBarry Smith   PetscInt          *owner, *start, *sizes, nsends, nreceives;
2691447629fSBarry Smith   PetscInt           nmax, count, *sindices, *rindices, idx, lastidx;
2701447629fSBarry Smith   PetscInt          *owners = aomems->map->range;
2711447629fSBarry Smith   MPI_Request       *send_waits, *recv_waits;
2721447629fSBarry Smith   MPI_Status         recv_status;
273*6497c311SBarry Smith   PetscMPIInt        widx;
2741447629fSBarry Smith   PetscInt          *rbuf;
2751447629fSBarry Smith   PetscInt           n = napp, ip, ia;
2761447629fSBarry Smith   MPI_Status        *send_status;
277*6497c311SBarry Smith   PetscCount         nindices;
2781447629fSBarry Smith 
2791447629fSBarry Smith   PetscFunctionBegin;
2809566063dSJacob Faibussowitsch   PetscCall(PetscArrayzero(aomap_loc, n_local));
2811447629fSBarry Smith 
2829566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
2839566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
2841447629fSBarry Smith 
2851447629fSBarry Smith   /*  first count number of contributors (of from_array[]) to each processor */
2869566063dSJacob Faibussowitsch   PetscCall(PetscCalloc1(2 * size, &sizes));
2879566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(n, &owner));
2881447629fSBarry Smith 
2891447629fSBarry Smith   j       = 0;
2901447629fSBarry Smith   lastidx = -1;
2911447629fSBarry Smith   for (i = 0; i < n; i++) {
2921447629fSBarry Smith     /* if indices are NOT locally sorted, need to start search at the beginning */
2931447629fSBarry Smith     if (lastidx > (idx = from_array[i])) j = 0;
2941447629fSBarry Smith     lastidx = idx;
2951447629fSBarry Smith     for (; j < size; j++) {
2961447629fSBarry Smith       if (idx >= owners[j] && idx < owners[j + 1]) {
29776ec1555SBarry Smith         sizes[2 * j] += 2;    /* num of indices to be sent - in pairs (ip,ia) */
29876ec1555SBarry Smith         sizes[2 * j + 1] = 1; /* send to proc[j] */
2991447629fSBarry Smith         owner[i]         = j;
3001447629fSBarry Smith         break;
3011447629fSBarry Smith       }
3021447629fSBarry Smith     }
3031447629fSBarry Smith   }
30476ec1555SBarry Smith   sizes[2 * rank] = sizes[2 * rank + 1] = 0; /* do not receive from self! */
3051447629fSBarry Smith   nsends                                = 0;
30676ec1555SBarry Smith   for (i = 0; i < size; i++) nsends += sizes[2 * i + 1];
3071447629fSBarry Smith 
3081447629fSBarry Smith   /* inform other processors of number of messages and max length*/
3099566063dSJacob Faibussowitsch   PetscCall(PetscMaxSum(comm, sizes, &nmax, &nreceives));
3101447629fSBarry Smith 
3111447629fSBarry Smith   /* allocate arrays */
3129566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)ao, &tag));
3139566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(nreceives * nmax, &rindices, nreceives, &recv_waits));
3149566063dSJacob Faibussowitsch   PetscCall(PetscMalloc3(2 * n, &sindices, nsends, &send_waits, nsends, &send_status));
3159566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(size, &start));
3161447629fSBarry Smith 
3171447629fSBarry Smith   /* post receives: */
318*6497c311SBarry Smith   for (i = 0; i < nreceives; i++) PetscCallMPI(MPIU_Irecv(rindices + nmax * i, nmax, MPIU_INT, MPI_ANY_SOURCE, tag, comm, recv_waits + i));
3191447629fSBarry Smith 
3201447629fSBarry Smith   /* do sends:
3211447629fSBarry Smith       1) starts[i] gives the starting index in svalues for stuff going to
3221447629fSBarry Smith          the ith processor
3231447629fSBarry Smith   */
3241447629fSBarry Smith   start[0] = 0;
32576ec1555SBarry Smith   for (i = 1; i < size; i++) start[i] = start[i - 1] + sizes[2 * i - 2];
3261447629fSBarry Smith   for (i = 0; i < n; i++) {
3271447629fSBarry Smith     j = owner[i];
3281447629fSBarry Smith     if (j != rank) {
3291447629fSBarry Smith       ip                   = from_array[i];
3301447629fSBarry Smith       ia                   = to_array[i];
3311447629fSBarry Smith       sindices[start[j]++] = ip;
3321447629fSBarry Smith       sindices[start[j]++] = ia;
3331447629fSBarry Smith     } else { /* compute my own map */
3341447629fSBarry Smith       ip            = from_array[i] - owners[rank];
3351447629fSBarry Smith       ia            = to_array[i];
3361447629fSBarry Smith       aomap_loc[ip] = ia;
3371447629fSBarry Smith     }
3381447629fSBarry Smith   }
3391447629fSBarry Smith 
3401447629fSBarry Smith   start[0] = 0;
341*6497c311SBarry Smith   for (PetscMPIInt i = 1; i < size; i++) start[i] = start[i - 1] + sizes[2 * i - 2];
342*6497c311SBarry Smith   count = 0;
343*6497c311SBarry Smith   for (PetscMPIInt i = 0; i < size; i++) {
34476ec1555SBarry Smith     if (sizes[2 * i + 1]) {
345*6497c311SBarry Smith       PetscCallMPI(MPIU_Isend(sindices + start[i], sizes[2 * i], MPIU_INT, i, tag, comm, send_waits + count));
3461447629fSBarry Smith       count++;
3471447629fSBarry Smith     }
3481447629fSBarry Smith   }
34908401ef6SPierre Jolivet   PetscCheck(nsends == count, comm, PETSC_ERR_SUP, "nsends %" PetscInt_FMT " != count %" PetscInt_FMT, nsends, count);
3501447629fSBarry Smith 
3511447629fSBarry Smith   /* wait on sends */
352*6497c311SBarry Smith   if (nsends) PetscCallMPI(MPI_Waitall((PetscMPIInt)nsends, send_waits, send_status));
3531447629fSBarry Smith 
3541447629fSBarry Smith   /* recvs */
3551447629fSBarry Smith   count = 0;
3561447629fSBarry Smith   for (j = nreceives; j > 0; j--) {
357*6497c311SBarry Smith     PetscCallMPI(MPI_Waitany((PetscMPIInt)nreceives, recv_waits, &widx, &recv_status));
358*6497c311SBarry Smith     PetscCallMPI(MPIU_Get_count(&recv_status, MPIU_INT, &nindices));
3591447629fSBarry Smith     rbuf = rindices + nmax * widx; /* global index */
3601447629fSBarry Smith 
3611447629fSBarry Smith     /* compute local mapping */
3621447629fSBarry Smith     for (i = 0; i < nindices; i += 2) {       /* pack aomap_loc */
3631447629fSBarry Smith       ip            = rbuf[i] - owners[rank]; /* local index */
3641447629fSBarry Smith       ia            = rbuf[i + 1];
3651447629fSBarry Smith       aomap_loc[ip] = ia;
3661447629fSBarry Smith     }
3671447629fSBarry Smith     count++;
3681447629fSBarry Smith   }
3691447629fSBarry Smith 
3709566063dSJacob Faibussowitsch   PetscCall(PetscFree(start));
3719566063dSJacob Faibussowitsch   PetscCall(PetscFree3(sindices, send_waits, send_status));
3729566063dSJacob Faibussowitsch   PetscCall(PetscFree2(rindices, recv_waits));
3739566063dSJacob Faibussowitsch   PetscCall(PetscFree(owner));
3749566063dSJacob Faibussowitsch   PetscCall(PetscFree(sizes));
3753ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3761447629fSBarry Smith }
3771447629fSBarry Smith 
378da8c939bSJacob Faibussowitsch PETSC_INTERN PetscErrorCode AOCreate_MemoryScalable(AO ao)
379d71ae5a4SJacob Faibussowitsch {
3801447629fSBarry Smith   IS                 isapp = ao->isapp, ispetsc = ao->ispetsc;
3811447629fSBarry Smith   const PetscInt    *mypetsc, *myapp;
3821447629fSBarry Smith   PetscInt           napp, n_local, N, i, start, *petsc, *lens, *disp;
3831447629fSBarry Smith   MPI_Comm           comm;
3841447629fSBarry Smith   AO_MemoryScalable *aomems;
3851447629fSBarry Smith   PetscLayout        map;
3861447629fSBarry Smith   PetscMPIInt        size, rank;
3871447629fSBarry Smith 
3881447629fSBarry Smith   PetscFunctionBegin;
38928b400f6SJacob Faibussowitsch   PetscCheck(isapp, PetscObjectComm((PetscObject)ao), PETSC_ERR_ARG_WRONGSTATE, "AOSetIS() must be called before AOSetType()");
3901447629fSBarry Smith   /* create special struct aomems */
3914dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&aomems));
3921447629fSBarry Smith   ao->data   = (void *)aomems;
393aea10558SJacob Faibussowitsch   ao->ops[0] = AOOps_MemoryScalable;
3949566063dSJacob Faibussowitsch   PetscCall(PetscObjectChangeTypeName((PetscObject)ao, AOMEMORYSCALABLE));
3951447629fSBarry Smith 
3961447629fSBarry Smith   /* transmit all local lengths of isapp to all processors */
3979566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)isapp, &comm));
3989566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
3999566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
4009566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(size, &lens, size, &disp));
4019566063dSJacob Faibussowitsch   PetscCall(ISGetLocalSize(isapp, &napp));
4029566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allgather(&napp, 1, MPIU_INT, lens, 1, MPIU_INT, comm));
4031447629fSBarry Smith 
4041447629fSBarry Smith   N = 0;
4051447629fSBarry Smith   for (i = 0; i < size; i++) {
4061447629fSBarry Smith     disp[i] = N;
4071447629fSBarry Smith     N += lens[i];
4081447629fSBarry Smith   }
4091447629fSBarry Smith 
4101447629fSBarry Smith   /* If ispetsc is 0 then use "natural" numbering */
4111447629fSBarry Smith   if (napp) {
4121447629fSBarry Smith     if (!ispetsc) {
4131447629fSBarry Smith       start = disp[rank];
4149566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(napp + 1, &petsc));
4151447629fSBarry Smith       for (i = 0; i < napp; i++) petsc[i] = start + i;
4161447629fSBarry Smith     } else {
4179566063dSJacob Faibussowitsch       PetscCall(ISGetIndices(ispetsc, &mypetsc));
4181447629fSBarry Smith       petsc = (PetscInt *)mypetsc;
4191447629fSBarry Smith     }
4204a2f8832SBarry Smith   } else {
4214a2f8832SBarry Smith     petsc = NULL;
4221447629fSBarry Smith   }
4231447629fSBarry Smith 
4241447629fSBarry Smith   /* create a map with global size N - used to determine the local sizes of ao - shall we use local napp instead of N? */
4259566063dSJacob Faibussowitsch   PetscCall(PetscLayoutCreate(comm, &map));
4261447629fSBarry Smith   map->bs = 1;
4271447629fSBarry Smith   map->N  = N;
4289566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(map));
4291447629fSBarry Smith 
4301447629fSBarry Smith   ao->N       = N;
4311447629fSBarry Smith   ao->n       = map->n;
4321447629fSBarry Smith   aomems->map = map;
4331447629fSBarry Smith 
4341447629fSBarry Smith   /* create distributed indices app_loc: petsc->app and petsc_loc: app->petsc */
4351447629fSBarry Smith   n_local = map->n;
4369566063dSJacob Faibussowitsch   PetscCall(PetscCalloc2(n_local, &aomems->app_loc, n_local, &aomems->petsc_loc));
4379566063dSJacob Faibussowitsch   PetscCall(ISGetIndices(isapp, &myapp));
4381447629fSBarry Smith 
4399566063dSJacob Faibussowitsch   PetscCall(AOCreateMemoryScalable_private(comm, napp, petsc, myapp, ao, aomems->app_loc));
4409566063dSJacob Faibussowitsch   PetscCall(AOCreateMemoryScalable_private(comm, napp, myapp, petsc, ao, aomems->petsc_loc));
4411447629fSBarry Smith 
4429566063dSJacob Faibussowitsch   PetscCall(ISRestoreIndices(isapp, &myapp));
4431447629fSBarry Smith   if (napp) {
4441447629fSBarry Smith     if (ispetsc) {
4459566063dSJacob Faibussowitsch       PetscCall(ISRestoreIndices(ispetsc, &mypetsc));
4461447629fSBarry Smith     } else {
4479566063dSJacob Faibussowitsch       PetscCall(PetscFree(petsc));
4481447629fSBarry Smith     }
4491447629fSBarry Smith   }
4509566063dSJacob Faibussowitsch   PetscCall(PetscFree2(lens, disp));
4513ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4521447629fSBarry Smith }
4531447629fSBarry Smith 
454cc4c1da9SBarry Smith /*@
4551447629fSBarry Smith   AOCreateMemoryScalable - Creates a memory scalable application ordering using two integer arrays.
4561447629fSBarry Smith 
457d083f849SBarry Smith   Collective
4581447629fSBarry Smith 
4591447629fSBarry Smith   Input Parameters:
460cab54364SBarry Smith + comm    - MPI communicator that is to share the `AO`
461d38ec673SBarry Smith . napp    - size of `myapp` and `mypetsc`
4621447629fSBarry Smith . myapp   - integer array that defines an ordering
463b6971eaeSBarry Smith - mypetsc - integer array that defines another ordering (may be `NULL` to indicate the natural ordering, that is 0,1,2,3,...)
4641447629fSBarry Smith 
4651447629fSBarry Smith   Output Parameter:
4661447629fSBarry Smith . aoout - the new application ordering
4671447629fSBarry Smith 
4681447629fSBarry Smith   Level: beginner
4691447629fSBarry Smith 
470cab54364SBarry Smith   Note:
471b6971eaeSBarry Smith   The arrays `myapp` and `mypetsc` must contain the all the integers 0 to `napp`-1 with no duplicates; that is there cannot be any "holes"
472cab54364SBarry Smith   in the indices. Use `AOCreateMapping()` or `AOCreateMappingIS()` if you wish to have "holes" in the indices.
473cab54364SBarry Smith   Comparing with `AOCreateBasic()`, this routine trades memory with message communication.
4741447629fSBarry Smith 
475cab54364SBarry Smith .seealso: [](sec_ao), [](sec_scatter), `AO`, `AOCreateMemoryScalableIS()`, `AODestroy()`, `AOPetscToApplication()`, `AOApplicationToPetsc()`
4761447629fSBarry Smith @*/
477d71ae5a4SJacob Faibussowitsch PetscErrorCode AOCreateMemoryScalable(MPI_Comm comm, PetscInt napp, const PetscInt myapp[], const PetscInt mypetsc[], AO *aoout)
478d71ae5a4SJacob Faibussowitsch {
4791447629fSBarry Smith   IS              isapp, ispetsc;
4801447629fSBarry Smith   const PetscInt *app = myapp, *petsc = mypetsc;
4811447629fSBarry Smith 
4821447629fSBarry Smith   PetscFunctionBegin;
4839566063dSJacob Faibussowitsch   PetscCall(ISCreateGeneral(comm, napp, app, PETSC_USE_POINTER, &isapp));
4841447629fSBarry Smith   if (mypetsc) {
4859566063dSJacob Faibussowitsch     PetscCall(ISCreateGeneral(comm, napp, petsc, PETSC_USE_POINTER, &ispetsc));
4861447629fSBarry Smith   } else {
4871447629fSBarry Smith     ispetsc = NULL;
4881447629fSBarry Smith   }
4899566063dSJacob Faibussowitsch   PetscCall(AOCreateMemoryScalableIS(isapp, ispetsc, aoout));
4909566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isapp));
49148a46eb9SPierre Jolivet   if (mypetsc) PetscCall(ISDestroy(&ispetsc));
4923ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4931447629fSBarry Smith }
4941447629fSBarry Smith 
4955d83a8b1SBarry Smith /*@
4961447629fSBarry Smith   AOCreateMemoryScalableIS - Creates a memory scalable application ordering using two index sets.
4971447629fSBarry Smith 
498c3339decSBarry Smith   Collective
4991447629fSBarry Smith 
5001447629fSBarry Smith   Input Parameters:
5011447629fSBarry Smith + isapp   - index set that defines an ordering
502b6971eaeSBarry Smith - ispetsc - index set that defines another ordering (may be `NULL` to use the natural ordering)
5031447629fSBarry Smith 
5041447629fSBarry Smith   Output Parameter:
5051447629fSBarry Smith . aoout - the new application ordering
5061447629fSBarry Smith 
5071447629fSBarry Smith   Level: beginner
5081447629fSBarry Smith 
50995452b02SPatrick Sanan   Notes:
510b6971eaeSBarry Smith   The index sets `isapp` and `ispetsc` must contain the all the integers 0 to napp-1 (where napp is the length of the index sets) with no duplicates;
5111447629fSBarry Smith   that is there cannot be any "holes".
512cab54364SBarry Smith 
513cab54364SBarry Smith   Comparing with `AOCreateBasicIS()`, this routine trades memory with message communication.
514cab54364SBarry Smith 
515cab54364SBarry Smith .seealso: [](sec_ao), [](sec_scatter), `AO`, `AOCreateBasicIS()`, `AOCreateMemoryScalable()`, `AODestroy()`
5161447629fSBarry Smith @*/
517d71ae5a4SJacob Faibussowitsch PetscErrorCode AOCreateMemoryScalableIS(IS isapp, IS ispetsc, AO *aoout)
518d71ae5a4SJacob Faibussowitsch {
5191447629fSBarry Smith   MPI_Comm comm;
5201447629fSBarry Smith   AO       ao;
5211447629fSBarry Smith 
5221447629fSBarry Smith   PetscFunctionBegin;
5239566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)isapp, &comm));
5249566063dSJacob Faibussowitsch   PetscCall(AOCreate(comm, &ao));
5259566063dSJacob Faibussowitsch   PetscCall(AOSetIS(ao, isapp, ispetsc));
5269566063dSJacob Faibussowitsch   PetscCall(AOSetType(ao, AOMEMORYSCALABLE));
5279566063dSJacob Faibussowitsch   PetscCall(AOViewFromOptions(ao, NULL, "-ao_view"));
5281447629fSBarry Smith   *aoout = ao;
5293ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5301447629fSBarry Smith }
531