Actual source code: mmaij.c


  2: /*
  3:    Support for the parallel AIJ matrix vector multiply
  4: */
  5: #include <../src/mat/impls/aij/mpi/mpiaij.h>
  6: #include <petsc/private/vecimpl.h>
  7: #include <petsc/private/isimpl.h>

  9: PetscErrorCode MatSetUpMultiply_MPIAIJ(Mat mat)
 10: {
 11:   Mat_MPIAIJ *aij = (Mat_MPIAIJ *)mat->data;
 12:   Mat_SeqAIJ *B   = (Mat_SeqAIJ *)(aij->B->data);
 13:   PetscInt    i, j, *aj = B->j, *garray;
 14:   PetscInt    ec = 0; /* Number of nonzero external columns */
 15:   IS          from, to;
 16:   Vec         gvec;
 17: #if defined(PETSC_USE_CTABLE)
 18:   PetscTable         gid1_lid1;
 19:   PetscTablePosition tpos;
 20:   PetscInt           gid, lid;
 21: #else
 22:   PetscInt N = mat->cmap->N, *indices;
 23: #endif

 25:   if (!aij->garray) {
 27: #if defined(PETSC_USE_CTABLE)
 28:     /* use a table */
 29:     PetscTableCreate(aij->B->rmap->n, mat->cmap->N + 1, &gid1_lid1);
 30:     for (i = 0; i < aij->B->rmap->n; i++) {
 31:       for (j = 0; j < B->ilen[i]; j++) {
 32:         PetscInt data, gid1 = aj[B->i[i] + j] + 1;
 33:         PetscTableFind(gid1_lid1, gid1, &data);
 34:         if (!data) {
 35:           /* one based table */
 36:           PetscTableAdd(gid1_lid1, gid1, ++ec, INSERT_VALUES);
 37:         }
 38:       }
 39:     }
 40:     /* form array of columns we need */
 41:     PetscMalloc1(ec, &garray);
 42:     PetscTableGetHeadPosition(gid1_lid1, &tpos);
 43:     while (tpos) {
 44:       PetscTableGetNext(gid1_lid1, &tpos, &gid, &lid);
 45:       gid--;
 46:       lid--;
 47:       garray[lid] = gid;
 48:     }
 49:     PetscSortInt(ec, garray); /* sort, and rebuild */
 50:     PetscTableRemoveAll(gid1_lid1);
 51:     for (i = 0; i < ec; i++) PetscTableAdd(gid1_lid1, garray[i] + 1, i + 1, INSERT_VALUES);
 52:     /* compact out the extra columns in B */
 53:     for (i = 0; i < aij->B->rmap->n; i++) {
 54:       for (j = 0; j < B->ilen[i]; j++) {
 55:         PetscInt gid1 = aj[B->i[i] + j] + 1;
 56:         PetscTableFind(gid1_lid1, gid1, &lid);
 57:         lid--;
 58:         aj[B->i[i] + j] = lid;
 59:       }
 60:     }
 61:     PetscLayoutDestroy(&aij->B->cmap);
 62:     PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)aij->B), ec, ec, 1, &aij->B->cmap);
 63:     PetscTableDestroy(&gid1_lid1);
 64: #else
 65:     /* Make an array as long as the number of columns */
 66:     /* mark those columns that are in aij->B */
 67:     PetscCalloc1(N, &indices);
 68:     for (i = 0; i < aij->B->rmap->n; i++) {
 69:       for (j = 0; j < B->ilen[i]; j++) {
 70:         if (!indices[aj[B->i[i] + j]]) ec++;
 71:         indices[aj[B->i[i] + j]] = 1;
 72:       }
 73:     }

 75:     /* form array of columns we need */
 76:     PetscMalloc1(ec, &garray);
 77:     ec = 0;
 78:     for (i = 0; i < N; i++) {
 79:       if (indices[i]) garray[ec++] = i;
 80:     }

 82:     /* make indices now point into garray */
 83:     for (i = 0; i < ec; i++) indices[garray[i]] = i;

 85:     /* compact out the extra columns in B */
 86:     for (i = 0; i < aij->B->rmap->n; i++) {
 87:       for (j = 0; j < B->ilen[i]; j++) aj[B->i[i] + j] = indices[aj[B->i[i] + j]];
 88:     }
 89:     PetscLayoutDestroy(&aij->B->cmap);
 90:     PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)aij->B), ec, ec, 1, &aij->B->cmap);
 91:     PetscFree(indices);
 92: #endif
 93:   } else {
 94:     garray = aij->garray;
 95:   }

 97:   if (!aij->lvec) {
 99:     MatCreateVecs(aij->B, &aij->lvec, NULL);
100:   }
101:   VecGetSize(aij->lvec, &ec);

103:   /* create two temporary Index sets for build scatter gather */
104:   ISCreateGeneral(PETSC_COMM_SELF, ec, garray, PETSC_COPY_VALUES, &from);
105:   ISCreateStride(PETSC_COMM_SELF, ec, 0, 1, &to);

107:   /* create temporary global vector to generate scatter context */
108:   /* This does not allocate the array's memory so is efficient */
109:   VecCreateMPIWithArray(PetscObjectComm((PetscObject)mat), 1, mat->cmap->n, mat->cmap->N, NULL, &gvec);

111:   /* generate the scatter context */
112:   VecScatterDestroy(&aij->Mvctx);
113:   VecScatterCreate(gvec, from, aij->lvec, to, &aij->Mvctx);
114:   VecScatterViewFromOptions(aij->Mvctx, (PetscObject)mat, "-matmult_vecscatter_view");
115:   aij->garray = garray;

117:   ISDestroy(&from);
118:   ISDestroy(&to);
119:   VecDestroy(&gvec);
120:   return 0;
121: }

123: /* Disassemble the off-diag portion of the MPIAIJXxx matrix.
124:    In other words, change the B from reduced format using local col ids
125:    to expanded format using global col ids, which will make it easier to
126:    insert new nonzeros (with global col ids) into the matrix.
127:    The off-diag B determines communication in the matrix vector multiply.
128: */
129: PetscErrorCode MatDisAssemble_MPIAIJ(Mat A)
130: {
131:   Mat_MPIAIJ        *aij  = (Mat_MPIAIJ *)A->data;
132:   Mat                B    = aij->B, Bnew;
133:   Mat_SeqAIJ        *Baij = (Mat_SeqAIJ *)B->data;
134:   PetscInt           i, j, m = B->rmap->n, n = A->cmap->N, col, ct = 0, *garray = aij->garray, *nz;
135:   PetscScalar        v;
136:   const PetscScalar *ba;

138:   /* free stuff related to matrix-vec multiply */
139:   VecDestroy(&aij->lvec);
140:   if (aij->colmap) {
141: #if defined(PETSC_USE_CTABLE)
142:     PetscTableDestroy(&aij->colmap);
143: #else
144:     PetscFree(aij->colmap);
145: #endif
146:   }

148:   /* make sure that B is assembled so we can access its values */
149:   MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY);
150:   MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY);

152:   /* invent new B and copy stuff over */
153:   PetscMalloc1(m + 1, &nz);
154:   for (i = 0; i < m; i++) nz[i] = Baij->i[i + 1] - Baij->i[i];
155:   MatCreate(PETSC_COMM_SELF, &Bnew);
156:   MatSetSizes(Bnew, m, n, m, n); /* Bnew now uses A->cmap->N as its col size */
157:   MatSetBlockSizesFromMats(Bnew, A, A);
158:   MatSetType(Bnew, ((PetscObject)B)->type_name);
159:   MatSeqAIJSetPreallocation(Bnew, 0, nz);

161:   if (Baij->nonew >= 0) { /* Inherit insertion error options (if positive). */
162:     ((Mat_SeqAIJ *)Bnew->data)->nonew = Baij->nonew;
163:   }

165:   /*
166:    Ensure that B's nonzerostate is monotonically increasing.
167:    Or should this follow the MatSetValues() loop to preserve B's nonzerstate across a MatDisAssemble() call?
168:    */
169:   Bnew->nonzerostate = B->nonzerostate;

171:   PetscFree(nz);
172:   MatSeqAIJGetArrayRead(B, &ba);
173:   for (i = 0; i < m; i++) {
174:     for (j = Baij->i[i]; j < Baij->i[i + 1]; j++) {
175:       col = garray[Baij->j[ct]];
176:       v   = ba[ct++];
177:       MatSetValues(Bnew, 1, &i, 1, &col, &v, B->insertmode);
178:     }
179:   }
180:   MatSeqAIJRestoreArrayRead(B, &ba);

182:   PetscFree(aij->garray);
183:   MatDestroy(&B);

185:   aij->B           = Bnew;
186:   A->was_assembled = PETSC_FALSE;
187:   A->assembled     = PETSC_FALSE;
188:   return 0;
189: }

191: /*      ugly stuff added for Glenn someday we should fix this up */

193: static PetscInt *auglyrmapd = NULL, *auglyrmapo = NULL; /* mapping from the local ordering to the "diagonal" and "off-diagonal" parts of the local matrix */
194: static Vec       auglydd = NULL, auglyoo = NULL;        /* work vectors used to scale the two parts of the local matrix */

196: PetscErrorCode MatMPIAIJDiagonalScaleLocalSetUp(Mat inA, Vec scale)
197: {
198:   Mat_MPIAIJ *ina = (Mat_MPIAIJ *)inA->data; /*access private part of matrix */
199:   PetscInt    i, n, nt, cstart, cend, no, *garray = ina->garray, *lindices;
200:   PetscInt   *r_rmapd, *r_rmapo;

202:   MatGetOwnershipRange(inA, &cstart, &cend);
203:   MatGetSize(ina->A, NULL, &n);
204:   PetscCalloc1(inA->rmap->mapping->n + 1, &r_rmapd);
205:   nt = 0;
206:   for (i = 0; i < inA->rmap->mapping->n; i++) {
207:     if (inA->rmap->mapping->indices[i] >= cstart && inA->rmap->mapping->indices[i] < cend) {
208:       nt++;
209:       r_rmapd[i] = inA->rmap->mapping->indices[i] + 1;
210:     }
211:   }
213:   PetscMalloc1(n + 1, &auglyrmapd);
214:   for (i = 0; i < inA->rmap->mapping->n; i++) {
215:     if (r_rmapd[i]) auglyrmapd[(r_rmapd[i] - 1) - cstart] = i;
216:   }
217:   PetscFree(r_rmapd);
218:   VecCreateSeq(PETSC_COMM_SELF, n, &auglydd);

220:   PetscCalloc1(inA->cmap->N + 1, &lindices);
221:   for (i = 0; i < ina->B->cmap->n; i++) lindices[garray[i]] = i + 1;
222:   no = inA->rmap->mapping->n - nt;
223:   PetscCalloc1(inA->rmap->mapping->n + 1, &r_rmapo);
224:   nt = 0;
225:   for (i = 0; i < inA->rmap->mapping->n; i++) {
226:     if (lindices[inA->rmap->mapping->indices[i]]) {
227:       nt++;
228:       r_rmapo[i] = lindices[inA->rmap->mapping->indices[i]];
229:     }
230:   }
232:   PetscFree(lindices);
233:   PetscMalloc1(nt + 1, &auglyrmapo);
234:   for (i = 0; i < inA->rmap->mapping->n; i++) {
235:     if (r_rmapo[i]) auglyrmapo[(r_rmapo[i] - 1)] = i;
236:   }
237:   PetscFree(r_rmapo);
238:   VecCreateSeq(PETSC_COMM_SELF, nt, &auglyoo);
239:   return 0;
240: }

242: PetscErrorCode MatMPIAIJDiagonalScaleLocal(Mat A, Vec scale)
243: {
244:   /* This routine should really be abandoned as it duplicates MatDiagonalScaleLocal */

246:   PetscTryMethod(A, "MatDiagonalScaleLocal_C", (Mat, Vec), (A, scale));
247:   return 0;
248: }

250: PetscErrorCode MatDiagonalScaleLocal_MPIAIJ(Mat A, Vec scale)
251: {
252:   Mat_MPIAIJ        *a = (Mat_MPIAIJ *)A->data; /*access private part of matrix */
253:   PetscInt           n, i;
254:   PetscScalar       *d, *o;
255:   const PetscScalar *s;

257:   if (!auglyrmapd) MatMPIAIJDiagonalScaleLocalSetUp(A, scale);

259:   VecGetArrayRead(scale, &s);

261:   VecGetLocalSize(auglydd, &n);
262:   VecGetArray(auglydd, &d);
263:   for (i = 0; i < n; i++) { d[i] = s[auglyrmapd[i]]; /* copy "diagonal" (true local) portion of scale into dd vector */ }
264:   VecRestoreArray(auglydd, &d);
265:   /* column scale "diagonal" portion of local matrix */
266:   MatDiagonalScale(a->A, NULL, auglydd);

268:   VecGetLocalSize(auglyoo, &n);
269:   VecGetArray(auglyoo, &o);
270:   for (i = 0; i < n; i++) { o[i] = s[auglyrmapo[i]]; /* copy "off-diagonal" portion of scale into oo vector */ }
271:   VecRestoreArrayRead(scale, &s);
272:   VecRestoreArray(auglyoo, &o);
273:   /* column scale "off-diagonal" portion of local matrix */
274:   MatDiagonalScale(a->B, NULL, auglyoo);
275:   return 0;
276: }