xref: /petsc/src/ksp/pc/impls/redistribute/redistribute.c (revision bcda9346efad4e5ba2d553af84eb238771ba1e25)
1df826632SBarry Smith /*
2df826632SBarry Smith   This file defines a "solve the problem redistributely on each subgroup of processor" preconditioner.
3df826632SBarry Smith */
4af0996ceSBarry Smith #include <petsc/private/pcimpl.h> /*I "petscksp.h" I*/
5c6db04a5SJed Brown #include <petscksp.h>
6df826632SBarry Smith 
769ff637eSBarry Smith typedef struct _PC_FieldSplitLink *PC_FieldSplitLink;
869ff637eSBarry Smith struct _PC_FieldSplitLink {
969ff637eSBarry Smith   char             *splitname;
1069ff637eSBarry Smith   IS                is;
1169ff637eSBarry Smith   PC_FieldSplitLink next, previous;
1269ff637eSBarry Smith };
1369ff637eSBarry Smith 
14df826632SBarry Smith typedef struct {
15df826632SBarry Smith   KSP          ksp;
16df826632SBarry Smith   Vec          x, b;
17df826632SBarry Smith   VecScatter   scatter;
18911f9fe8SBarry Smith   IS           is;
19181dd334SBarry Smith   PetscInt     dcnt, *drows; /* these are the local rows that have only diagonal entry */
20181dd334SBarry Smith   PetscScalar *diag;
21181dd334SBarry Smith   Vec          work;
2287b47708SBarry Smith   PetscBool    zerodiag;
2369ff637eSBarry Smith 
2469ff637eSBarry Smith   PetscInt          nsplits;
2569ff637eSBarry Smith   PC_FieldSplitLink splitlinks;
26df826632SBarry Smith } PC_Redistribute;
27df826632SBarry Smith 
PCFieldSplitSetIS_Redistribute(PC pc,const char splitname[],IS is)2869ff637eSBarry Smith static PetscErrorCode PCFieldSplitSetIS_Redistribute(PC pc, const char splitname[], IS is)
2969ff637eSBarry Smith {
3069ff637eSBarry Smith   PC_Redistribute   *red  = (PC_Redistribute *)pc->data;
3169ff637eSBarry Smith   PC_FieldSplitLink *next = &red->splitlinks;
3269ff637eSBarry Smith 
3369ff637eSBarry Smith   PetscFunctionBegin;
3469ff637eSBarry Smith   while (*next) next = &(*next)->next;
3569ff637eSBarry Smith   PetscCall(PetscNew(next));
3669ff637eSBarry Smith   if (splitname) {
3769ff637eSBarry Smith     PetscCall(PetscStrallocpy(splitname, &(*next)->splitname));
3869ff637eSBarry Smith   } else {
3969ff637eSBarry Smith     PetscCall(PetscMalloc1(8, &(*next)->splitname));
4069ff637eSBarry Smith     PetscCall(PetscSNPrintf((*next)->splitname, 7, "%" PetscInt_FMT, red->nsplits++));
4169ff637eSBarry Smith   }
4269ff637eSBarry Smith   PetscCall(PetscObjectReference((PetscObject)is));
4369ff637eSBarry Smith   PetscCall(ISDestroy(&(*next)->is));
4469ff637eSBarry Smith   (*next)->is = is;
4569ff637eSBarry Smith   PetscFunctionReturn(PETSC_SUCCESS);
4669ff637eSBarry Smith }
4769ff637eSBarry Smith 
PCView_Redistribute(PC pc,PetscViewer viewer)48d71ae5a4SJacob Faibussowitsch static PetscErrorCode PCView_Redistribute(PC pc, PetscViewer viewer)
49d71ae5a4SJacob Faibussowitsch {
50df826632SBarry Smith   PC_Redistribute *red = (PC_Redistribute *)pc->data;
51*9f196a02SMartin Diehl   PetscBool        isascii, isstring;
52181dd334SBarry Smith   PetscInt         ncnt, N;
53df826632SBarry Smith 
54df826632SBarry Smith   PetscFunctionBegin;
55*9f196a02SMartin Diehl   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
569566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSTRING, &isstring));
57*9f196a02SMartin Diehl   if (isascii) {
58462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(&red->dcnt, &ncnt, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)pc)));
599566063dSJacob Faibussowitsch     PetscCall(MatGetSize(pc->pmat, &N, NULL));
60835f2295SStefano Zampini     PetscCall(PetscViewerASCIIPrintf(viewer, "    Number rows eliminated %" PetscInt_FMT " Percentage rows eliminated %g\n", ncnt, (double)(100 * ((PetscReal)ncnt) / ((PetscReal)N))));
619566063dSJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(viewer, "  Redistribute preconditioner: \n"));
629566063dSJacob Faibussowitsch     PetscCall(KSPView(red->ksp, viewer));
63df826632SBarry Smith   } else if (isstring) {
649566063dSJacob Faibussowitsch     PetscCall(PetscViewerStringSPrintf(viewer, " Redistribute preconditioner"));
659566063dSJacob Faibussowitsch     PetscCall(KSPView(red->ksp, viewer));
6611aeaf0aSBarry Smith   }
673ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
68df826632SBarry Smith }
69df826632SBarry Smith 
PCSetUp_Redistribute(PC pc)70d71ae5a4SJacob Faibussowitsch static PetscErrorCode PCSetUp_Redistribute(PC pc)
71d71ae5a4SJacob Faibussowitsch {
72df826632SBarry Smith   PC_Redistribute         *red = (PC_Redistribute *)pc->data;
73df826632SBarry Smith   MPI_Comm                 comm;
746497c311SBarry Smith   PetscInt                 rstart, rend, nrstart, nrend, nz, cnt, *rows, ncnt, dcnt, *drows;
7526283091SBarry Smith   PetscLayout              map, nmap;
76ec4bef21SJose E. Roman   PetscMPIInt              size, tag, n;
77ec4bef21SJose E. Roman   PETSC_UNUSED PetscMPIInt imdex;
780298fd71SBarry Smith   PetscInt                *source = NULL;
796497c311SBarry Smith   PetscMPIInt             *sizes  = NULL, nrecvs, nsends;
806497c311SBarry Smith   PetscInt                 j;
810298fd71SBarry Smith   PetscInt                *owner = NULL, *starts = NULL, count, slen;
82e8dd6687SHong Zhang   PetscInt                *rvalues, *svalues, recvtotal;
83df826632SBarry Smith   PetscMPIInt             *onodes1, *olengths1;
840298fd71SBarry Smith   MPI_Request             *send_waits = NULL, *recv_waits = NULL;
85df826632SBarry Smith   MPI_Status               recv_status, *send_status;
86181dd334SBarry Smith   Vec                      tvec, diag;
87ca320bd4SBarry Smith   Mat                      tmat;
88003249c0SBarry Smith   const PetscScalar       *d, *values;
89003249c0SBarry Smith   const PetscInt          *cols;
9069ff637eSBarry Smith   PC_FieldSplitLink       *next = &red->splitlinks;
91df826632SBarry Smith 
92df826632SBarry Smith   PetscFunctionBegin;
93dc9360f3SBarry Smith   if (pc->setupcalled) {
9469ff637eSBarry Smith     PetscCheck(pc->flag == SAME_NONZERO_PATTERN, PetscObjectComm((PetscObject)pc), PETSC_ERR_SUP, "PC is not supported for a change in the nonzero structure of the matrix");
959566063dSJacob Faibussowitsch     PetscCall(KSPGetOperators(red->ksp, NULL, &tmat));
969566063dSJacob Faibussowitsch     PetscCall(MatCreateSubMatrix(pc->pmat, red->is, red->is, MAT_REUSE_MATRIX, &tmat));
979566063dSJacob Faibussowitsch     PetscCall(KSPSetOperators(red->ksp, tmat, tmat));
98dc9360f3SBarry Smith   } else {
99b862ddfaSBarry Smith     PetscInt  NN;
10069ff637eSBarry Smith     PC        ipc;
1010cd8b6e2SStefano Zampini     PetscBool fptr;
102b862ddfaSBarry Smith 
1039566063dSJacob Faibussowitsch     PetscCall(PetscObjectGetComm((PetscObject)pc, &comm));
1049566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_size(comm, &size));
1059566063dSJacob Faibussowitsch     PetscCall(PetscObjectGetNewTag((PetscObject)pc, &tag));
106df826632SBarry Smith 
107ca320bd4SBarry Smith     /* count non-diagonal rows on process */
1089566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(pc->mat, &rstart, &rend));
109ca320bd4SBarry Smith     cnt = 0;
1106497c311SBarry Smith     for (PetscInt i = rstart; i < rend; i++) {
1119566063dSJacob Faibussowitsch       PetscCall(MatGetRow(pc->mat, i, &nz, &cols, &values));
112003249c0SBarry Smith       for (PetscInt j = 0; j < nz; j++) {
113003249c0SBarry Smith         if (values[j] != 0 && cols[j] != i) {
114003249c0SBarry Smith           cnt++;
115003249c0SBarry Smith           break;
116003249c0SBarry Smith         }
117003249c0SBarry Smith       }
1189566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(pc->mat, i, &nz, &cols, &values));
119df826632SBarry Smith     }
1209566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(cnt, &rows));
1219566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(rend - rstart - cnt, &drows));
122ca320bd4SBarry Smith 
123ca320bd4SBarry Smith     /* list non-diagonal rows on process */
1249371c9d4SSatish Balay     cnt  = 0;
1259371c9d4SSatish Balay     dcnt = 0;
1266497c311SBarry Smith     for (PetscInt i = rstart; i < rend; i++) {
127003249c0SBarry Smith       PetscBool diagonly = PETSC_TRUE;
1289566063dSJacob Faibussowitsch       PetscCall(MatGetRow(pc->mat, i, &nz, &cols, &values));
129003249c0SBarry Smith       for (PetscInt j = 0; j < nz; j++) {
130003249c0SBarry Smith         if (values[j] != 0 && cols[j] != i) {
131003249c0SBarry Smith           diagonly = PETSC_FALSE;
132003249c0SBarry Smith           break;
133003249c0SBarry Smith         }
134003249c0SBarry Smith       }
135003249c0SBarry Smith       if (!diagonly) rows[cnt++] = i;
136181dd334SBarry Smith       else drows[dcnt++] = i - rstart;
1379566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(pc->mat, i, &nz, &cols, &values));
138df826632SBarry Smith     }
139ca320bd4SBarry Smith 
14026283091SBarry Smith     /* create PetscLayout for non-diagonal rows on each process */
1419566063dSJacob Faibussowitsch     PetscCall(PetscLayoutCreate(comm, &map));
1429566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetLocalSize(map, cnt));
1439566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetBlockSize(map, 1));
1449566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(map));
14569ff637eSBarry Smith     nrstart = map->rstart;
14669ff637eSBarry Smith     nrend   = map->rend;
147df826632SBarry Smith 
14826283091SBarry Smith     /* create PetscLayout for load-balanced non-diagonal rows on each process */
1499566063dSJacob Faibussowitsch     PetscCall(PetscLayoutCreate(comm, &nmap));
150462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(&cnt, &ncnt, 1, MPIU_INT, MPI_SUM, comm));
1519566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetSize(nmap, ncnt));
1529566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetBlockSize(nmap, 1));
1539566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(nmap));
154df826632SBarry Smith 
1559566063dSJacob Faibussowitsch     PetscCall(MatGetSize(pc->pmat, &NN, NULL));
15657508eceSPierre Jolivet     PetscCall(PetscInfo(pc, "Number of diagonal rows eliminated %" PetscInt_FMT ", percentage eliminated %g\n", NN - ncnt, (double)((PetscReal)(NN - ncnt) / (PetscReal)NN)));
157003249c0SBarry Smith 
158003249c0SBarry Smith     if (size > 1) {
159ca320bd4SBarry Smith       /*
16069ff637eSBarry Smith         the following block of code assumes MPI can send messages to self, which is not supported for MPI-uni hence we need to handle
16169ff637eSBarry Smith         the size 1 case as a special case
16269ff637eSBarry Smith 
163ca320bd4SBarry Smith        this code is taken from VecScatterCreate_PtoS()
164ca320bd4SBarry Smith        Determines what rows need to be moved where to
165ca320bd4SBarry Smith        load balance the non-diagonal rows
166ca320bd4SBarry Smith        */
167df826632SBarry Smith       /*  count number of contributors to each processor */
1689566063dSJacob Faibussowitsch       PetscCall(PetscMalloc2(size, &sizes, cnt, &owner));
1699566063dSJacob Faibussowitsch       PetscCall(PetscArrayzero(sizes, size));
170df826632SBarry Smith       j      = 0;
171df826632SBarry Smith       nsends = 0;
1726497c311SBarry Smith       for (PetscInt i = nrstart; i < nrend; i++) {
173df826632SBarry Smith         if (i < nmap->range[j]) j = 0;
174df826632SBarry Smith         for (; j < size; j++) {
175df826632SBarry Smith           if (i < nmap->range[j + 1]) {
17676ec1555SBarry Smith             if (!sizes[j]++) nsends++;
17769ff637eSBarry Smith             owner[i - nrstart] = j;
178df826632SBarry Smith             break;
179df826632SBarry Smith           }
180df826632SBarry Smith         }
181df826632SBarry Smith       }
182df826632SBarry Smith       /* inform other processors of number of messages and max length*/
1839566063dSJacob Faibussowitsch       PetscCall(PetscGatherNumberOfMessages(comm, NULL, sizes, &nrecvs));
1849566063dSJacob Faibussowitsch       PetscCall(PetscGatherMessageLengths(comm, nsends, nrecvs, sizes, &onodes1, &olengths1));
1859566063dSJacob Faibussowitsch       PetscCall(PetscSortMPIIntWithArray(nrecvs, onodes1, olengths1));
1869371c9d4SSatish Balay       recvtotal = 0;
1876497c311SBarry Smith       for (PetscMPIInt i = 0; i < nrecvs; i++) recvtotal += olengths1[i];
188df826632SBarry Smith 
189df826632SBarry Smith       /* post receives:  rvalues - rows I will own; count - nu */
1909566063dSJacob Faibussowitsch       PetscCall(PetscMalloc3(recvtotal, &rvalues, nrecvs, &source, nrecvs, &recv_waits));
191df826632SBarry Smith       count = 0;
1926497c311SBarry Smith       for (PetscMPIInt i = 0; i < nrecvs; i++) {
19357508eceSPierre Jolivet         PetscCallMPI(MPIU_Irecv(rvalues + count, olengths1[i], MPIU_INT, onodes1[i], tag, comm, recv_waits + i));
194df826632SBarry Smith         count += olengths1[i];
195df826632SBarry Smith       }
196df826632SBarry Smith 
197df826632SBarry Smith       /* do sends:
198df826632SBarry Smith        1) starts[i] gives the starting index in svalues for stuff going to
199df826632SBarry Smith        the ith processor
200df826632SBarry Smith        */
2019566063dSJacob Faibussowitsch       PetscCall(PetscMalloc3(cnt, &svalues, nsends, &send_waits, size, &starts));
202df826632SBarry Smith       starts[0] = 0;
2036497c311SBarry Smith       for (PetscMPIInt i = 1; i < size; i++) starts[i] = starts[i - 1] + sizes[i - 1];
2046497c311SBarry Smith       for (PetscInt i = 0; i < cnt; i++) svalues[starts[owner[i]]++] = rows[i];
2056497c311SBarry Smith       for (PetscInt i = 0; i < cnt; i++) rows[i] = rows[i] - nrstart;
206181dd334SBarry Smith       red->drows = drows;
207181dd334SBarry Smith       red->dcnt  = dcnt;
2089566063dSJacob Faibussowitsch       PetscCall(PetscFree(rows));
209181dd334SBarry Smith 
210df826632SBarry Smith       starts[0] = 0;
2116497c311SBarry Smith       for (PetscMPIInt i = 1; i < size; i++) starts[i] = starts[i - 1] + sizes[i - 1];
212df826632SBarry Smith       count = 0;
2136497c311SBarry Smith       for (PetscMPIInt i = 0; i < size; i++) {
2146497c311SBarry Smith         if (sizes[i]) PetscCallMPI(MPIU_Isend(svalues + starts[i], sizes[i], MPIU_INT, i, tag, comm, send_waits + count++));
215df826632SBarry Smith       }
216df826632SBarry Smith 
217df826632SBarry Smith       /*  wait on receives */
218df826632SBarry Smith       count = nrecvs;
219df826632SBarry Smith       slen  = 0;
220df826632SBarry Smith       while (count) {
2219566063dSJacob Faibussowitsch         PetscCallMPI(MPI_Waitany(nrecvs, recv_waits, &imdex, &recv_status));
222df826632SBarry Smith         /* unpack receives into our local space */
2239566063dSJacob Faibussowitsch         PetscCallMPI(MPI_Get_count(&recv_status, MPIU_INT, &n));
224df826632SBarry Smith         slen += n;
225df826632SBarry Smith         count--;
226df826632SBarry Smith       }
22763a3b9bcSJacob Faibussowitsch       PetscCheck(slen == recvtotal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Total message lengths %" PetscInt_FMT " not expected %" PetscInt_FMT, slen, recvtotal);
2289566063dSJacob Faibussowitsch       PetscCall(ISCreateGeneral(comm, slen, rvalues, PETSC_COPY_VALUES, &red->is));
229911f9fe8SBarry Smith 
230003249c0SBarry Smith       /* free all work space */
2319566063dSJacob Faibussowitsch       PetscCall(PetscFree(olengths1));
2329566063dSJacob Faibussowitsch       PetscCall(PetscFree(onodes1));
2339566063dSJacob Faibussowitsch       PetscCall(PetscFree3(rvalues, source, recv_waits));
2349566063dSJacob Faibussowitsch       PetscCall(PetscFree2(sizes, owner));
235ca320bd4SBarry Smith       if (nsends) { /* wait on sends */
2369566063dSJacob Faibussowitsch         PetscCall(PetscMalloc1(nsends, &send_status));
2379566063dSJacob Faibussowitsch         PetscCallMPI(MPI_Waitall(nsends, send_waits, send_status));
2389566063dSJacob Faibussowitsch         PetscCall(PetscFree(send_status));
239df826632SBarry Smith       }
2409566063dSJacob Faibussowitsch       PetscCall(PetscFree3(svalues, send_waits, starts));
241003249c0SBarry Smith     } else {
2429566063dSJacob Faibussowitsch       PetscCall(ISCreateGeneral(comm, cnt, rows, PETSC_OWN_POINTER, &red->is));
243003249c0SBarry Smith       red->drows = drows;
244003249c0SBarry Smith       red->dcnt  = dcnt;
245003249c0SBarry Smith       slen       = cnt;
246003249c0SBarry Smith     }
2479566063dSJacob Faibussowitsch     PetscCall(PetscLayoutDestroy(&map));
248df826632SBarry Smith 
2499566063dSJacob Faibussowitsch     PetscCall(VecCreateMPI(comm, slen, PETSC_DETERMINE, &red->b));
2509566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(red->b, &red->x));
2519566063dSJacob Faibussowitsch     PetscCall(MatCreateVecs(pc->pmat, &tvec, NULL));
2529566063dSJacob Faibussowitsch     PetscCall(VecScatterCreate(tvec, red->is, red->b, NULL, &red->scatter));
25369ff637eSBarry Smith 
25469ff637eSBarry Smith     /* Map the PCFIELDSPLIT fields to redistributed KSP */
25569ff637eSBarry Smith     PetscCall(KSPGetPC(red->ksp, &ipc));
2560cd8b6e2SStefano Zampini     PetscCall(PetscObjectHasFunction((PetscObject)ipc, "PCFieldSplitSetIS_C", &fptr));
25769ff637eSBarry Smith     if (fptr && *next) {
25869ff637eSBarry Smith       PetscScalar       *atvec;
25969ff637eSBarry Smith       const PetscScalar *ab;
26069ff637eSBarry Smith       PetscInt           primes[] = {2, 3, 5, 7, 11, 13, 17, 19};
26169ff637eSBarry Smith       PetscInt           cnt      = 0;
26269ff637eSBarry Smith 
26369ff637eSBarry Smith       PetscCheck(red->nsplits <= (PetscInt)PETSC_STATIC_ARRAY_LENGTH(primes), PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "No support for this many fields");
26469ff637eSBarry Smith       PetscCall(VecSet(tvec, 1.0));
26569ff637eSBarry Smith       PetscCall(VecGetArray(tvec, &atvec));
26669ff637eSBarry Smith 
26769ff637eSBarry Smith       while (*next) {
26869ff637eSBarry Smith         const PetscInt *indices;
26969ff637eSBarry Smith         PetscInt        n;
27069ff637eSBarry Smith 
27169ff637eSBarry Smith         PetscCall(ISGetIndices((*next)->is, &indices));
27269ff637eSBarry Smith         PetscCall(ISGetLocalSize((*next)->is, &n));
27369ff637eSBarry Smith         for (PetscInt i = 0; i < n; i++) atvec[indices[i] - rstart] *= primes[cnt];
27469ff637eSBarry Smith         PetscCall(ISRestoreIndices((*next)->is, &indices));
27569ff637eSBarry Smith         cnt++;
27669ff637eSBarry Smith         next = &(*next)->next;
27769ff637eSBarry Smith       }
27869ff637eSBarry Smith       PetscCall(VecRestoreArray(tvec, &atvec));
27969ff637eSBarry Smith       PetscCall(VecScatterBegin(red->scatter, tvec, red->b, INSERT_VALUES, SCATTER_FORWARD));
28069ff637eSBarry Smith       PetscCall(VecScatterEnd(red->scatter, tvec, red->b, INSERT_VALUES, SCATTER_FORWARD));
28169ff637eSBarry Smith       cnt = 0;
28269ff637eSBarry Smith       PetscCall(VecGetArrayRead(red->b, &ab));
28369ff637eSBarry Smith       next = &red->splitlinks;
28469ff637eSBarry Smith       while (*next) {
28569ff637eSBarry Smith         PetscInt  n = 0;
28669ff637eSBarry Smith         PetscInt *indices;
28769ff637eSBarry Smith         IS        ris;
28869ff637eSBarry Smith 
28969ff637eSBarry Smith         for (PetscInt i = 0; i < nmap->rend - nmap->rstart; i++) {
29069ff637eSBarry Smith           if (!(((PetscInt)PetscRealPart(ab[i])) % primes[cnt])) n++;
29169ff637eSBarry Smith         }
29269ff637eSBarry Smith         PetscCall(PetscMalloc1(n, &indices));
29369ff637eSBarry Smith         n = 0;
29469ff637eSBarry Smith         for (PetscInt i = 0; i < nmap->rend - nmap->rstart; i++) {
29569ff637eSBarry Smith           if (!(((PetscInt)PetscRealPart(ab[i])) % primes[cnt])) indices[n++] = i + nmap->rstart;
29669ff637eSBarry Smith         }
29769ff637eSBarry Smith         PetscCall(ISCreateGeneral(comm, n, indices, PETSC_OWN_POINTER, &ris));
29869ff637eSBarry Smith         PetscCall(PCFieldSplitSetIS(ipc, (*next)->splitname, ris));
29969ff637eSBarry Smith 
30069ff637eSBarry Smith         PetscCall(ISDestroy(&ris));
30169ff637eSBarry Smith         cnt++;
30269ff637eSBarry Smith         next = &(*next)->next;
30369ff637eSBarry Smith       }
30469ff637eSBarry Smith       PetscCall(VecRestoreArrayRead(red->b, &ab));
30569ff637eSBarry Smith     }
3069566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&tvec));
3079566063dSJacob Faibussowitsch     PetscCall(MatCreateSubMatrix(pc->pmat, red->is, red->is, MAT_INITIAL_MATRIX, &tmat));
3089566063dSJacob Faibussowitsch     PetscCall(KSPSetOperators(red->ksp, tmat, tmat));
3099566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&tmat));
31069ff637eSBarry Smith     PetscCall(PetscLayoutDestroy(&nmap));
3111d805cfdSBarry Smith   }
312181dd334SBarry Smith 
313181dd334SBarry Smith   /* get diagonal portion of matrix */
3149566063dSJacob Faibussowitsch   PetscCall(PetscFree(red->diag));
3159566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(red->dcnt, &red->diag));
3169566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(pc->pmat, &diag, NULL));
3179566063dSJacob Faibussowitsch   PetscCall(MatGetDiagonal(pc->pmat, diag));
3189566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(diag, &d));
3196497c311SBarry Smith   for (PetscInt i = 0; i < red->dcnt; i++) {
32087b47708SBarry Smith     if (d[red->drows[i]] != 0) red->diag[i] = 1.0 / d[red->drows[i]];
32187b47708SBarry Smith     else {
32287b47708SBarry Smith       red->zerodiag = PETSC_TRUE;
32387b47708SBarry Smith       red->diag[i]  = 0.0;
32487b47708SBarry Smith     }
32587b47708SBarry Smith   }
3269566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(diag, &d));
3279566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&diag));
3289566063dSJacob Faibussowitsch   PetscCall(KSPSetUp(red->ksp));
3293ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
330df826632SBarry Smith }
331df826632SBarry Smith 
PCApply_Redistribute(PC pc,Vec b,Vec x)332d71ae5a4SJacob Faibussowitsch static PetscErrorCode PCApply_Redistribute(PC pc, Vec b, Vec x)
333d71ae5a4SJacob Faibussowitsch {
334df826632SBarry Smith   PC_Redistribute   *red   = (PC_Redistribute *)pc->data;
335181dd334SBarry Smith   PetscInt           dcnt  = red->dcnt, i;
336181dd334SBarry Smith   const PetscInt    *drows = red->drows;
337181dd334SBarry Smith   PetscScalar       *xwork;
338181dd334SBarry Smith   const PetscScalar *bwork, *diag = red->diag;
3393ae97395SBarry Smith   PetscBool          nonzero_guess;
340df826632SBarry Smith 
341df826632SBarry Smith   PetscFunctionBegin;
34248a46eb9SPierre Jolivet   if (!red->work) PetscCall(VecDuplicate(b, &red->work));
3433ae97395SBarry Smith   PetscCall(KSPGetInitialGuessNonzero(red->ksp, &nonzero_guess));
3443ae97395SBarry Smith   if (nonzero_guess) {
3453ae97395SBarry Smith     PetscCall(VecScatterBegin(red->scatter, x, red->x, INSERT_VALUES, SCATTER_FORWARD));
3463ae97395SBarry Smith     PetscCall(VecScatterEnd(red->scatter, x, red->x, INSERT_VALUES, SCATTER_FORWARD));
3473ae97395SBarry Smith   }
3483ae97395SBarry Smith 
349181dd334SBarry Smith   /* compute the rows of solution that have diagonal entries only */
3509566063dSJacob Faibussowitsch   PetscCall(VecSet(x, 0.0)); /* x = diag(A)^{-1} b */
3519566063dSJacob Faibussowitsch   PetscCall(VecGetArray(x, &xwork));
3529566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(b, &bwork));
35387b47708SBarry Smith   if (red->zerodiag) {
35487b47708SBarry Smith     for (i = 0; i < dcnt; i++) {
35587b47708SBarry Smith       if (diag[i] == 0.0 && bwork[drows[i]] != 0.0) {
356dd8e379bSPierre Jolivet         PetscCheck(!pc->erroriffailure, PETSC_COMM_SELF, PETSC_ERR_CONV_FAILED, "Linear system is inconsistent, zero matrix row but nonzero right-hand side");
357dd8e379bSPierre Jolivet         PetscCall(PetscInfo(pc, "Linear system is inconsistent, zero matrix row but nonzero right-hand side\n"));
35887b47708SBarry Smith         pc->failedreasonrank = PC_INCONSISTENT_RHS;
35987b47708SBarry Smith       }
36087b47708SBarry Smith     }
361f480ea8aSBarry Smith     PetscCall(VecFlag(x, pc->failedreasonrank == PC_INCONSISTENT_RHS));
36287b47708SBarry Smith   }
3632fa5cd67SKarl Rupp   for (i = 0; i < dcnt; i++) xwork[drows[i]] = diag[i] * bwork[drows[i]];
3649566063dSJacob Faibussowitsch   PetscCall(PetscLogFlops(dcnt));
3659566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(red->work, &xwork));
3669566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(b, &bwork));
367dd8e379bSPierre Jolivet   /* update the right-hand side for the reduced system with diagonal rows (and corresponding columns) removed */
3689566063dSJacob Faibussowitsch   PetscCall(MatMult(pc->pmat, x, red->work));
3699566063dSJacob Faibussowitsch   PetscCall(VecAYPX(red->work, -1.0, b)); /* red->work = b - A x */
370181dd334SBarry Smith 
3719566063dSJacob Faibussowitsch   PetscCall(VecScatterBegin(red->scatter, red->work, red->b, INSERT_VALUES, SCATTER_FORWARD));
3729566063dSJacob Faibussowitsch   PetscCall(VecScatterEnd(red->scatter, red->work, red->b, INSERT_VALUES, SCATTER_FORWARD));
3739566063dSJacob Faibussowitsch   PetscCall(KSPSolve(red->ksp, red->b, red->x));
3749566063dSJacob Faibussowitsch   PetscCall(KSPCheckSolve(red->ksp, pc, red->x));
3759566063dSJacob Faibussowitsch   PetscCall(VecScatterBegin(red->scatter, red->x, x, INSERT_VALUES, SCATTER_REVERSE));
3769566063dSJacob Faibussowitsch   PetscCall(VecScatterEnd(red->scatter, red->x, x, INSERT_VALUES, SCATTER_REVERSE));
3773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
378df826632SBarry Smith }
379df826632SBarry Smith 
PCApplyTranspose_Redistribute(PC pc,Vec b,Vec x)38057cdce21SPierre Jolivet static PetscErrorCode PCApplyTranspose_Redistribute(PC pc, Vec b, Vec x)
38157cdce21SPierre Jolivet {
38257cdce21SPierre Jolivet   PC_Redistribute   *red   = (PC_Redistribute *)pc->data;
38357cdce21SPierre Jolivet   PetscInt           dcnt  = red->dcnt, i;
38457cdce21SPierre Jolivet   const PetscInt    *drows = red->drows;
38557cdce21SPierre Jolivet   PetscScalar       *xwork;
38657cdce21SPierre Jolivet   const PetscScalar *bwork, *diag = red->diag;
38757cdce21SPierre Jolivet   PetscBool          set, flg     = PETSC_FALSE, nonzero_guess;
38857cdce21SPierre Jolivet 
38957cdce21SPierre Jolivet   PetscFunctionBegin;
39057cdce21SPierre Jolivet   PetscCall(MatIsStructurallySymmetricKnown(pc->pmat, &set, &flg));
39157cdce21SPierre Jolivet   PetscCheck(set || flg, PetscObjectComm((PetscObject)pc), PETSC_ERR_SUP, "PCApplyTranspose() not implemented for structurally unsymmetric Mat");
39257cdce21SPierre Jolivet   if (!red->work) PetscCall(VecDuplicate(b, &red->work));
39357cdce21SPierre Jolivet   PetscCall(KSPGetInitialGuessNonzero(red->ksp, &nonzero_guess));
39457cdce21SPierre Jolivet   if (nonzero_guess) {
39557cdce21SPierre Jolivet     PetscCall(VecScatterBegin(red->scatter, x, red->x, INSERT_VALUES, SCATTER_FORWARD));
39657cdce21SPierre Jolivet     PetscCall(VecScatterEnd(red->scatter, x, red->x, INSERT_VALUES, SCATTER_FORWARD));
39757cdce21SPierre Jolivet   }
39857cdce21SPierre Jolivet 
39957cdce21SPierre Jolivet   /* compute the rows of solution that have diagonal entries only */
40057cdce21SPierre Jolivet   PetscCall(VecSet(x, 0.0)); /* x = diag(A)^{-1} b */
40157cdce21SPierre Jolivet   PetscCall(VecGetArray(x, &xwork));
40257cdce21SPierre Jolivet   PetscCall(VecGetArrayRead(b, &bwork));
40357cdce21SPierre Jolivet   if (red->zerodiag) {
40457cdce21SPierre Jolivet     for (i = 0; i < dcnt; i++) {
40557cdce21SPierre Jolivet       if (diag[i] == 0.0 && bwork[drows[i]] != 0.0) {
40657cdce21SPierre Jolivet         PetscCheck(!pc->erroriffailure, PETSC_COMM_SELF, PETSC_ERR_CONV_FAILED, "Linear system is inconsistent, zero matrix row but nonzero right-hand side");
40757cdce21SPierre Jolivet         PetscCall(PetscInfo(pc, "Linear system is inconsistent, zero matrix row but nonzero right-hand side\n"));
40857cdce21SPierre Jolivet         pc->failedreasonrank = PC_INCONSISTENT_RHS;
40957cdce21SPierre Jolivet       }
41057cdce21SPierre Jolivet     }
411f480ea8aSBarry Smith     PetscCall(VecFlag(x, pc->failedreasonrank == PC_INCONSISTENT_RHS));
41257cdce21SPierre Jolivet   }
41357cdce21SPierre Jolivet   for (i = 0; i < dcnt; i++) xwork[drows[i]] = diag[i] * bwork[drows[i]];
41457cdce21SPierre Jolivet   PetscCall(PetscLogFlops(dcnt));
41557cdce21SPierre Jolivet   PetscCall(VecRestoreArray(red->work, &xwork));
41657cdce21SPierre Jolivet   PetscCall(VecRestoreArrayRead(b, &bwork));
41757cdce21SPierre Jolivet   /* update the right-hand side for the reduced system with diagonal rows (and corresponding columns) removed */
41857cdce21SPierre Jolivet   PetscCall(MatMultTranspose(pc->pmat, x, red->work));
41957cdce21SPierre Jolivet   PetscCall(VecAYPX(red->work, -1.0, b)); /* red->work = b - A^T x */
42057cdce21SPierre Jolivet 
42157cdce21SPierre Jolivet   PetscCall(VecScatterBegin(red->scatter, red->work, red->b, INSERT_VALUES, SCATTER_FORWARD));
42257cdce21SPierre Jolivet   PetscCall(VecScatterEnd(red->scatter, red->work, red->b, INSERT_VALUES, SCATTER_FORWARD));
42357cdce21SPierre Jolivet   PetscCall(KSPSolveTranspose(red->ksp, red->b, red->x));
42457cdce21SPierre Jolivet   PetscCall(KSPCheckSolve(red->ksp, pc, red->x));
42557cdce21SPierre Jolivet   PetscCall(VecScatterBegin(red->scatter, red->x, x, INSERT_VALUES, SCATTER_REVERSE));
42657cdce21SPierre Jolivet   PetscCall(VecScatterEnd(red->scatter, red->x, x, INSERT_VALUES, SCATTER_REVERSE));
42757cdce21SPierre Jolivet   PetscFunctionReturn(PETSC_SUCCESS);
42857cdce21SPierre Jolivet }
42957cdce21SPierre Jolivet 
PCDestroy_Redistribute(PC pc)430d71ae5a4SJacob Faibussowitsch static PetscErrorCode PCDestroy_Redistribute(PC pc)
431d71ae5a4SJacob Faibussowitsch {
432df826632SBarry Smith   PC_Redistribute  *red  = (PC_Redistribute *)pc->data;
43369ff637eSBarry Smith   PC_FieldSplitLink next = red->splitlinks;
434df826632SBarry Smith 
435df826632SBarry Smith   PetscFunctionBegin;
43669ff637eSBarry Smith   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCFieldSplitSetIS_C", NULL));
43769ff637eSBarry Smith 
43869ff637eSBarry Smith   while (next) {
43969ff637eSBarry Smith     PC_FieldSplitLink ilink;
44069ff637eSBarry Smith     PetscCall(PetscFree(next->splitname));
44169ff637eSBarry Smith     PetscCall(ISDestroy(&next->is));
44269ff637eSBarry Smith     ilink = next;
44369ff637eSBarry Smith     next  = next->next;
44469ff637eSBarry Smith     PetscCall(PetscFree(ilink));
44569ff637eSBarry Smith   }
4469566063dSJacob Faibussowitsch   PetscCall(VecScatterDestroy(&red->scatter));
4479566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&red->is));
4489566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&red->b));
4499566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&red->x));
4509566063dSJacob Faibussowitsch   PetscCall(KSPDestroy(&red->ksp));
4519566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&red->work));
4529566063dSJacob Faibussowitsch   PetscCall(PetscFree(red->drows));
4539566063dSJacob Faibussowitsch   PetscCall(PetscFree(red->diag));
4549566063dSJacob Faibussowitsch   PetscCall(PetscFree(pc->data));
4553ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
456df826632SBarry Smith }
457df826632SBarry Smith 
PCSetFromOptions_Redistribute(PC pc,PetscOptionItems PetscOptionsObject)458ce78bad3SBarry Smith static PetscErrorCode PCSetFromOptions_Redistribute(PC pc, PetscOptionItems PetscOptionsObject)
459d71ae5a4SJacob Faibussowitsch {
460df826632SBarry Smith   PC_Redistribute *red = (PC_Redistribute *)pc->data;
461df826632SBarry Smith 
462df826632SBarry Smith   PetscFunctionBegin;
4639566063dSJacob Faibussowitsch   PetscCall(KSPSetFromOptions(red->ksp));
4643ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
465df826632SBarry Smith }
466df826632SBarry Smith 
4675e7ef714SBarry Smith /*@
468f1580f4eSBarry Smith   PCRedistributeGetKSP - Gets the `KSP` created by the `PCREDISTRIBUTE`
4695e7ef714SBarry Smith 
4705e7ef714SBarry Smith   Not Collective
4715e7ef714SBarry Smith 
4725e7ef714SBarry Smith   Input Parameter:
4735e7ef714SBarry Smith . pc - the preconditioner context
4745e7ef714SBarry Smith 
4755e7ef714SBarry Smith   Output Parameter:
476f1580f4eSBarry Smith . innerksp - the inner `KSP`
4775e7ef714SBarry Smith 
4785e7ef714SBarry Smith   Level: advanced
4795e7ef714SBarry Smith 
480562efe2eSBarry Smith .seealso: [](ch_ksp), `KSP`, `PCREDISTRIBUTE`
4815e7ef714SBarry Smith @*/
PCRedistributeGetKSP(PC pc,KSP * innerksp)482d71ae5a4SJacob Faibussowitsch PetscErrorCode PCRedistributeGetKSP(PC pc, KSP *innerksp)
483d71ae5a4SJacob Faibussowitsch {
4845e7ef714SBarry Smith   PC_Redistribute *red = (PC_Redistribute *)pc->data;
4855e7ef714SBarry Smith 
4865e7ef714SBarry Smith   PetscFunctionBegin;
4875e7ef714SBarry Smith   PetscValidHeaderSpecific(pc, PC_CLASSID, 1);
4884f572ea9SToby Isaac   PetscAssertPointer(innerksp, 2);
4895e7ef714SBarry Smith   *innerksp = red->ksp;
4903ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4915e7ef714SBarry Smith }
4925e7ef714SBarry Smith 
493df826632SBarry Smith /*MC
4947caa67d9SBarry Smith      PCREDISTRIBUTE - Redistributes a matrix for load balancing, removing the rows (and the corresponding columns) that only have a diagonal entry and then
495f1580f4eSBarry Smith      applies a `KSP` to that new smaller matrix
496df826632SBarry Smith 
4977df6f684SBarry Smith      Level: intermediate
4987df6f684SBarry Smith 
49995452b02SPatrick Sanan      Notes:
500a76f6b59SBarry Smith      Options for the redistribute `KSP` and `PC` with the options database prefix `-redistribute_`
501f1580f4eSBarry Smith 
5027df6f684SBarry Smith      Usually run this with `-ksp_type preonly`
503181dd334SBarry Smith 
5047df6f684SBarry Smith      If you have used `MatZeroRows()` to eliminate (for example, Dirichlet) boundary conditions for a symmetric problem then you can use, for example, `-ksp_type preonly
5057df6f684SBarry Smith      -pc_type redistribute -redistribute_ksp_type cg -redistribute_pc_type bjacobi -redistribute_sub_pc_type icc` to take advantage of the symmetry.
506181dd334SBarry Smith 
507baca6076SPierre Jolivet      Supports the function `PCFieldSplitSetIS()`; pass the appropriate reduced field indices to an inner `PCFIELDSPLIT`, set with, for example
508a76f6b59SBarry Smith      `-ksp_type preonly -pc_type redistribute -redistribute_pc_type fieldsplit`. Does not support the `PCFIELDSPLIT` options database keys.
50969ff637eSBarry Smith 
5107df6f684SBarry Smith      This does NOT call a partitioner to reorder rows to lower communication; the ordering of the rows in the original matrix and redistributed matrix is the same. Rows are moved
5117df6f684SBarry Smith      between MPI processes inside the preconditioner to balance the number of rows on each process.
5122dfef595SBarry Smith 
513a76f6b59SBarry Smith      The matrix block information is lost with the possible removal of individual rows and columns of the matrix, thus the behavior of the preconditioner on the reduced
514a76f6b59SBarry Smith      system may be very different (worse) than running that preconditioner on the full system. This is specifically true for elasticity problems.
515a76f6b59SBarry Smith 
516f1580f4eSBarry Smith      Developer Note:
51795452b02SPatrick Sanan      Should add an option to this preconditioner to use a partitioner to redistribute the rows to lower communication.
5182dfef595SBarry Smith 
519562efe2eSBarry Smith .seealso: [](ch_ksp), `PCCreate()`, `PCSetType()`, `PCType`, `PCRedistributeGetKSP()`, `MatZeroRows()`, `PCFieldSplitSetIS()`, `PCFIELDSPLIT`
520df826632SBarry Smith M*/
521df826632SBarry Smith 
PCCreate_Redistribute(PC pc)522d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode PCCreate_Redistribute(PC pc)
523d71ae5a4SJacob Faibussowitsch {
524df826632SBarry Smith   PC_Redistribute *red;
525911f9fe8SBarry Smith   const char      *prefix;
526df826632SBarry Smith 
527df826632SBarry Smith   PetscFunctionBegin;
5284dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&red));
529df826632SBarry Smith   pc->data = (void *)red;
530df826632SBarry Smith 
531df826632SBarry Smith   pc->ops->apply          = PCApply_Redistribute;
53257cdce21SPierre Jolivet   pc->ops->applytranspose = PCApplyTranspose_Redistribute;
533df826632SBarry Smith   pc->ops->setup          = PCSetUp_Redistribute;
534df826632SBarry Smith   pc->ops->destroy        = PCDestroy_Redistribute;
535df826632SBarry Smith   pc->ops->setfromoptions = PCSetFromOptions_Redistribute;
536df826632SBarry Smith   pc->ops->view           = PCView_Redistribute;
537911f9fe8SBarry Smith 
5389566063dSJacob Faibussowitsch   PetscCall(KSPCreate(PetscObjectComm((PetscObject)pc), &red->ksp));
5393821be0aSBarry Smith   PetscCall(KSPSetNestLevel(red->ksp, pc->kspnestlevel));
5409566063dSJacob Faibussowitsch   PetscCall(KSPSetErrorIfNotConverged(red->ksp, pc->erroriffailure));
5419566063dSJacob Faibussowitsch   PetscCall(PetscObjectIncrementTabLevel((PetscObject)red->ksp, (PetscObject)pc, 1));
5429566063dSJacob Faibussowitsch   PetscCall(PCGetOptionsPrefix(pc, &prefix));
5439566063dSJacob Faibussowitsch   PetscCall(KSPSetOptionsPrefix(red->ksp, prefix));
5449566063dSJacob Faibussowitsch   PetscCall(KSPAppendOptionsPrefix(red->ksp, "redistribute_"));
54569ff637eSBarry Smith   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCFieldSplitSetIS_C", PCFieldSplitSetIS_Redistribute));
5463ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
547df826632SBarry Smith }
548