Actual source code: mmsbaij.c
1: /*
2: Support for the parallel SBAIJ matrix vector multiply
3: */
4: #include <../src/mat/impls/sbaij/mpi/mpisbaij.h>
6: PetscErrorCode MatSetUpMultiply_MPISBAIJ(Mat mat)
7: {
8: Mat_MPISBAIJ *sbaij = (Mat_MPISBAIJ *)mat->data;
9: Mat_SeqBAIJ *B = (Mat_SeqBAIJ *)sbaij->B->data;
10: PetscInt i, j, *aj = B->j, ec = 0, *garray, *sgarray;
11: PetscInt bs = mat->rmap->bs, *stmp, mbs = sbaij->mbs, vec_size, nt;
12: IS from, to;
13: Vec gvec;
14: PetscMPIInt rank = sbaij->rank;
15: PetscInt *owners = sbaij->rangebs, *ec_owner, k;
16: const PetscInt *sowners;
17: PetscScalar *ptr;
18: #if defined(PETSC_USE_CTABLE)
19: PetscHMapI gid1_lid1 = NULL; /* one-based gid to lid table */
20: PetscHashIter tpos;
21: PetscInt gid, lid;
22: #else
23: PetscInt Nbs = sbaij->Nbs;
24: PetscInt *indices;
25: #endif
27: PetscFunctionBegin;
28: #if defined(PETSC_USE_CTABLE)
29: PetscCall(PetscHMapICreateWithSize(mbs, &gid1_lid1));
30: for (i = 0; i < B->mbs; i++) {
31: for (j = 0; j < B->ilen[i]; j++) {
32: PetscInt data, gid1 = aj[B->i[i] + j] + 1;
33: PetscCall(PetscHMapIGetWithDefault(gid1_lid1, gid1, 0, &data));
34: if (!data) PetscCall(PetscHMapISet(gid1_lid1, gid1, ++ec));
35: }
36: }
37: /* form array of columns we need */
38: PetscCall(PetscMalloc1(ec, &garray));
39: PetscHashIterBegin(gid1_lid1, tpos);
40: while (!PetscHashIterAtEnd(gid1_lid1, tpos)) {
41: PetscHashIterGetKey(gid1_lid1, tpos, gid);
42: PetscHashIterGetVal(gid1_lid1, tpos, lid);
43: PetscHashIterNext(gid1_lid1, tpos);
44: gid--;
45: lid--;
46: garray[lid] = gid;
47: }
48: PetscCall(PetscSortInt(ec, garray));
49: PetscCall(PetscHMapIClear(gid1_lid1));
50: for (i = 0; i < ec; i++) PetscCall(PetscHMapISet(gid1_lid1, garray[i] + 1, i + 1));
51: /* compact out the extra columns in B */
52: for (i = 0; i < B->mbs; i++) {
53: for (j = 0; j < B->ilen[i]; j++) {
54: PetscInt gid1 = aj[B->i[i] + j] + 1;
55: PetscCall(PetscHMapIGetWithDefault(gid1_lid1, gid1, 0, &lid));
56: lid--;
57: aj[B->i[i] + j] = lid;
58: }
59: }
60: PetscCall(PetscHMapIDestroy(&gid1_lid1));
61: PetscCall(PetscMalloc2(2 * ec, &sgarray, ec, &ec_owner));
62: for (i = j = 0; i < ec; i++) {
63: while (garray[i] >= owners[j + 1]) j++;
64: ec_owner[i] = j;
65: }
66: #else
67: /* For the first stab we make an array as long as the number of columns */
68: /* mark those columns that are in sbaij->B */
69: PetscCall(PetscCalloc1(Nbs, &indices));
70: for (i = 0; i < mbs; i++) {
71: for (j = 0; j < B->ilen[i]; j++) {
72: if (!indices[aj[B->i[i] + j]]) ec++;
73: indices[aj[B->i[i] + j]] = 1;
74: }
75: }
77: /* form arrays of columns we need */
78: PetscCall(PetscMalloc1(ec, &garray));
79: PetscCall(PetscMalloc2(2 * ec, &sgarray, ec, &ec_owner));
81: ec = 0;
82: for (j = 0; j < sbaij->size; j++) {
83: for (i = owners[j]; i < owners[j + 1]; i++) {
84: if (indices[i]) {
85: garray[ec] = i;
86: ec_owner[ec] = j;
87: ec++;
88: }
89: }
90: }
92: /* make indices now point into garray */
93: for (i = 0; i < ec; i++) indices[garray[i]] = i;
95: /* compact out the extra columns in B */
96: for (i = 0; i < mbs; i++) {
97: for (j = 0; j < B->ilen[i]; j++) aj[B->i[i] + j] = indices[aj[B->i[i] + j]];
98: }
99: PetscCall(PetscFree(indices));
100: #endif
101: B->nbs = ec;
102: PetscCall(PetscLayoutDestroy(&sbaij->B->cmap));
103: PetscCall(PetscLayoutCreateFromSizes(PetscObjectComm((PetscObject)sbaij->B), ec * mat->rmap->bs, ec * mat->rmap->bs, mat->rmap->bs, &sbaij->B->cmap));
105: PetscCall(VecScatterDestroy(&sbaij->sMvctx));
106: /* create local vector that is used to scatter into */
107: PetscCall(VecCreateSeq(PETSC_COMM_SELF, ec * bs, &sbaij->lvec));
109: /* create two temporary index sets for building scatter-gather */
110: PetscCall(PetscMalloc1(2 * ec, &stmp));
111: PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, ec, garray, PETSC_COPY_VALUES, &from));
112: for (i = 0; i < ec; i++) stmp[i] = i;
113: PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, ec, stmp, PETSC_COPY_VALUES, &to));
115: /* generate the scatter context
116: -- Mvctx and lvec are not used by MatMult_MPISBAIJ(), but have other uses, such as in MatDiagonalScale_MPISBAIJ */
117: PetscCall(VecCreateMPIWithArray(PetscObjectComm((PetscObject)mat), 1, mat->cmap->n, mat->cmap->N, NULL, &gvec));
118: PetscCall(VecScatterCreate(gvec, from, sbaij->lvec, to, &sbaij->Mvctx));
119: PetscCall(VecDestroy(&gvec));
121: sbaij->garray = garray;
123: PetscCall(ISDestroy(&from));
124: PetscCall(ISDestroy(&to));
126: /* create parallel vector that is used by SBAIJ matrix to scatter from/into */
127: PetscCall(VecCreateMPI(PetscObjectComm((PetscObject)mat), (mbs + ec) * bs, PETSC_DETERMINE, &sbaij->slvec0));
128: PetscCall(VecDuplicate(sbaij->slvec0, &sbaij->slvec1));
129: PetscCall(VecGetSize(sbaij->slvec0, &vec_size));
131: PetscCall(VecGetOwnershipRanges(sbaij->slvec0, &sowners));
133: /* x index in the IS sfrom */
134: for (i = 0; i < ec; i++) {
135: j = ec_owner[i];
136: sgarray[i] = garray[i] + (sowners[j] / bs - owners[j]);
137: }
138: /* b index in the IS sfrom */
139: k = sowners[rank] / bs + mbs;
140: for (i = ec, j = 0; i < 2 * ec; i++, j++) sgarray[i] = k + j;
141: PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, 2 * ec, sgarray, PETSC_COPY_VALUES, &from));
143: /* x index in the IS sto */
144: k = sowners[rank] / bs + mbs;
145: for (i = 0; i < ec; i++) stmp[i] = (k + i);
146: /* b index in the IS sto */
147: for (i = ec; i < 2 * ec; i++) stmp[i] = sgarray[i - ec];
149: PetscCall(ISCreateBlock(PETSC_COMM_SELF, bs, 2 * ec, stmp, PETSC_COPY_VALUES, &to));
151: PetscCall(VecScatterCreate(sbaij->slvec0, from, sbaij->slvec1, to, &sbaij->sMvctx));
152: PetscCall(VecScatterViewFromOptions(sbaij->sMvctx, (PetscObject)mat, "-matmult_vecscatter_view"));
154: PetscCall(VecGetLocalSize(sbaij->slvec1, &nt));
155: PetscCall(VecGetArray(sbaij->slvec1, &ptr));
156: PetscCall(VecCreateSeqWithArray(PETSC_COMM_SELF, 1, bs * mbs, ptr, &sbaij->slvec1a));
157: PetscCall(VecCreateSeqWithArray(PETSC_COMM_SELF, 1, nt - bs * mbs, PetscSafePointerPlusOffset(ptr, bs * mbs), &sbaij->slvec1b));
158: PetscCall(VecRestoreArray(sbaij->slvec1, &ptr));
160: PetscCall(VecGetArray(sbaij->slvec0, &ptr));
161: PetscCall(VecCreateSeqWithArray(PETSC_COMM_SELF, 1, nt - bs * mbs, PetscSafePointerPlusOffset(ptr, bs * mbs), &sbaij->slvec0b));
162: PetscCall(VecRestoreArray(sbaij->slvec0, &ptr));
164: PetscCall(PetscFree(stmp));
166: PetscCall(ISDestroy(&from));
167: PetscCall(ISDestroy(&to));
168: PetscCall(PetscFree2(sgarray, ec_owner));
169: PetscFunctionReturn(PETSC_SUCCESS);
170: }
172: /*
173: Takes the local part of an already assembled MPISBAIJ matrix
174: and disassembles it. This is to allow new nonzeros into the matrix
175: that require more communication in the matrix vector multiply.
176: Thus certain data-structures must be rebuilt.
178: Kind of slow! But that's what application programmers get when
179: they are sloppy.
180: */
181: PetscErrorCode MatDisAssemble_MPISBAIJ(Mat A)
182: {
183: Mat_MPISBAIJ *baij = (Mat_MPISBAIJ *)A->data;
184: Mat B = baij->B, Bnew;
185: Mat_SeqBAIJ *Bbaij = (Mat_SeqBAIJ *)B->data;
186: PetscInt i, j, mbs = Bbaij->mbs, n = A->cmap->N, col, *garray = baij->garray;
187: PetscInt k, bs = A->rmap->bs, bs2 = baij->bs2, *rvals, *nz, m = A->rmap->n;
188: MatScalar *a = Bbaij->a;
189: PetscScalar *atmp;
190: #if defined(PETSC_USE_REAL_MAT_SINGLE)
191: PetscInt l;
192: #endif
194: PetscFunctionBegin;
195: #if defined(PETSC_USE_REAL_MAT_SINGLE)
196: PetscCall(PetscMalloc1(A->rmap->bs, &atmp));
197: #endif
198: /* free stuff related to matrix-vec multiply */
199: PetscCall(VecDestroy(&baij->lvec));
200: PetscCall(VecScatterDestroy(&baij->Mvctx));
202: PetscCall(VecDestroy(&baij->slvec0));
203: PetscCall(VecDestroy(&baij->slvec0b));
204: PetscCall(VecDestroy(&baij->slvec1));
205: PetscCall(VecDestroy(&baij->slvec1a));
206: PetscCall(VecDestroy(&baij->slvec1b));
208: if (baij->colmap) {
209: #if defined(PETSC_USE_CTABLE)
210: PetscCall(PetscHMapIDestroy(&baij->colmap));
211: #else
212: PetscCall(PetscFree(baij->colmap));
213: #endif
214: }
216: /* make sure that B is assembled so we can access its values */
217: PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
218: PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
220: /* invent new B and copy stuff over */
221: PetscCall(PetscMalloc1(mbs, &nz));
222: for (i = 0; i < mbs; i++) nz[i] = Bbaij->i[i + 1] - Bbaij->i[i];
223: PetscCall(MatCreate(PETSC_COMM_SELF, &Bnew));
224: PetscCall(MatSetSizes(Bnew, m, n, m, n));
225: PetscCall(MatSetType(Bnew, ((PetscObject)B)->type_name));
226: PetscCall(MatSeqBAIJSetPreallocation(Bnew, B->rmap->bs, 0, nz));
227: PetscCall(PetscFree(nz));
229: if (Bbaij->nonew >= 0) { /* Inherit insertion error options (if positive). */
230: ((Mat_SeqSBAIJ *)Bnew->data)->nonew = Bbaij->nonew;
231: }
233: /*
234: Ensure that B's nonzerostate is monotonically increasing.
235: Or should this follow the MatSetValues() loop to preserve B's nonzerstate across a MatDisAssemble() call?
236: */
237: Bnew->nonzerostate = B->nonzerostate;
238: PetscCall(PetscMalloc1(bs, &rvals));
239: for (i = 0; i < mbs; i++) {
240: rvals[0] = bs * i;
241: for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
242: for (j = Bbaij->i[i]; j < Bbaij->i[i + 1]; j++) {
243: col = garray[Bbaij->j[j]] * bs;
244: for (k = 0; k < bs; k++) {
245: #if defined(PETSC_USE_REAL_MAT_SINGLE)
246: for (l = 0; l < bs; l++) atmp[l] = a[j * bs2 + l];
247: #else
248: atmp = a + j * bs2 + k * bs;
249: #endif
250: PetscCall(MatSetValues_SeqSBAIJ(Bnew, bs, rvals, 1, &col, atmp, B->insertmode));
251: col++;
252: }
253: }
254: }
255: #if defined(PETSC_USE_REAL_MAT_SINGLE)
256: PetscCall(PetscFree(atmp));
257: #endif
258: PetscCall(PetscFree(baij->garray));
260: baij->garray = NULL;
262: PetscCall(PetscFree(rvals));
263: PetscCall(MatDestroy(&B));
265: baij->B = Bnew;
267: A->was_assembled = PETSC_FALSE;
268: PetscFunctionReturn(PETSC_SUCCESS);
269: }