Actual source code: mmbaij.c

      1: #define PETSCMAT_DLL

      3: /*
      4:    Support for the parallel BAIJ matrix vector multiply
      5: */
     6:  #include src/mat/impls/baij/mpi/mpibaij.h

      8: EXTERN PetscErrorCode MatSetValuesBlocked_SeqBAIJ(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode);

     12: PetscErrorCode MatSetUpMultiply_MPIBAIJ(Mat mat)
     13: {
     14:   Mat_MPIBAIJ    *baij = (Mat_MPIBAIJ*)mat->data;
     15:   Mat_SeqBAIJ    *B = (Mat_SeqBAIJ*)(baij->B->data);
     17:   PetscInt       i,j,*aj = B->j,ec = 0,*garray;
     18:   PetscInt       bs = mat->bs,*stmp;
     19:   IS             from,to;
     20:   Vec            gvec;
     21: #if defined (PETSC_USE_CTABLE)
     22:   PetscTable     gid1_lid1;
     23:   PetscTablePosition tpos;
     24:   PetscInt       gid,lid;
     25: #else
     26:   PetscInt       Nbs = baij->Nbs,*indices;
     27: #endif

     31: #if defined (PETSC_USE_CTABLE)
     32:   /* use a table - Mark Adams */
     33:   PetscTableCreate(B->mbs,&gid1_lid1);
     34:   for (i=0; i<B->mbs; i++) {
     35:     for (j=0; j<B->ilen[i]; j++) {
     36:       PetscInt data,gid1 = aj[B->i[i]+j] + 1;
     37:       PetscTableFind(gid1_lid1,gid1,&data);
     38:       if (!data) {
     39:         /* one based table */
     40:         PetscTableAdd(gid1_lid1,gid1,++ec);
     41:       }
     42:     }
     43:   }
     44:   /* form array of columns we need */
     45:   PetscMalloc((ec+1)*sizeof(PetscInt),&garray);
     46:   PetscTableGetHeadPosition(gid1_lid1,&tpos);
     47:   while (tpos) {
     48:     PetscTableGetNext(gid1_lid1,&tpos,&gid,&lid);
     49:     gid--; lid--;
     50:     garray[lid] = gid;
     51:   }
     52:   PetscSortInt(ec,garray);
     53:   PetscTableRemoveAll(gid1_lid1);
     54:   for (i=0; i<ec; i++) {
     55:     PetscTableAdd(gid1_lid1,garray[i]+1,i+1);
     56:   }
     57:   /* compact out the extra columns in B */
     58:   for (i=0; i<B->mbs; i++) {
     59:     for (j=0; j<B->ilen[i]; j++) {
     60:       PetscInt gid1 = aj[B->i[i] + j] + 1;
     61:       PetscTableFind(gid1_lid1,gid1,&lid);
     62:       lid --;
     63:       aj[B->i[i]+j] = lid;
     64:     }
     65:   }
     66:   B->nbs     = ec;
     67:   baij->B->n = ec*mat->bs;
     68:   PetscTableDelete(gid1_lid1);
     69:   /* Mark Adams */
     70: #else
     71:   /* Make an array as long as the number of columns */
     72:   /* mark those columns that are in baij->B */
     73:   PetscMalloc((Nbs+1)*sizeof(PetscInt),&indices);
     74:   PetscMemzero(indices,Nbs*sizeof(PetscInt));
     75:   for (i=0; i<B->mbs; i++) {
     76:     for (j=0; j<B->ilen[i]; j++) {
     77:       if (!indices[aj[B->i[i] + j]]) ec++;
     78:       indices[aj[B->i[i] + j]] = 1;
     79:     }
     80:   }

     82:   /* form array of columns we need */
     83:   PetscMalloc((ec+1)*sizeof(PetscInt),&garray);
     84:   ec = 0;
     85:   for (i=0; i<Nbs; i++) {
     86:     if (indices[i]) {
     87:       garray[ec++] = i;
     88:     }
     89:   }

     91:   /* make indices now point into garray */
     92:   for (i=0; i<ec; i++) {
     93:     indices[garray[i]] = i;
     94:   }

     96:   /* compact out the extra columns in B */
     97:   for (i=0; i<B->mbs; i++) {
     98:     for (j=0; j<B->ilen[i]; j++) {
     99:       aj[B->i[i] + j] = indices[aj[B->i[i] + j]];
    100:     }
    101:   }
    102:   B->nbs       = ec;
    103:   baij->B->n   = ec*mat->bs;
    104:   PetscFree(indices);
    105: #endif

    107:   /* create local vector that is used to scatter into */
    108:   VecCreateSeq(PETSC_COMM_SELF,ec*bs,&baij->lvec);

    110:   /* create two temporary index sets for building scatter-gather */
    111:   for (i=0; i<ec; i++) {
    112:     garray[i] = bs*garray[i];
    113:   }
    114:   ISCreateBlock(PETSC_COMM_SELF,bs,ec,garray,&from);
    115:   for (i=0; i<ec; i++) {
    116:     garray[i] = garray[i]/bs;
    117:   }

    119:   PetscMalloc((ec+1)*sizeof(PetscInt),&stmp);
    120:   for (i=0; i<ec; i++) { stmp[i] = bs*i; }
    121:   ISCreateBlock(PETSC_COMM_SELF,bs,ec,stmp,&to);
    122:   PetscFree(stmp);

    124:   /* create temporary global vector to generate scatter context */
    125:   /* this is inefficient, but otherwise we must do either
    126:      1) save garray until the first actual scatter when the vector is known or
    127:      2) have another way of generating a scatter context without a vector.*/
    128:   VecCreateMPI(mat->comm,mat->n,mat->N,&gvec);

    130:   /* gnerate the scatter context */
    131:   VecScatterCreate(gvec,from,baij->lvec,to,&baij->Mvctx);

    133:   /*
    134:       Post the receives for the first matrix vector product. We sync-chronize after
    135:     this on the chance that the user immediately calls MatMult() after assemblying
    136:     the matrix.
    137:   */
    138:   VecScatterPostRecvs(gvec,baij->lvec,INSERT_VALUES,SCATTER_FORWARD,baij->Mvctx);
    139:   MPI_Barrier(mat->comm);

    141:   PetscLogObjectParent(mat,baij->Mvctx);
    142:   PetscLogObjectParent(mat,baij->lvec);
    143:   PetscLogObjectParent(mat,from);
    144:   PetscLogObjectParent(mat,to);
    145:   baij->garray = garray;
    146:   PetscLogObjectMemory(mat,(ec+1)*sizeof(PetscInt));
    147:   ISDestroy(from);
    148:   ISDestroy(to);
    149:   VecDestroy(gvec);
    150:   return(0);
    151: }

    153: /*
    154:      Takes the local part of an already assembled MPIBAIJ matrix
    155:    and disassembles it. This is to allow new nonzeros into the matrix
    156:    that require more communication in the matrix vector multiply.
    157:    Thus certain data-structures must be rebuilt.

    159:    Kind of slow! But that's what application programmers get when
    160:    they are sloppy.
    161: */
    164: PetscErrorCode DisAssemble_MPIBAIJ(Mat A)
    165: {
    166:   Mat_MPIBAIJ    *baij = (Mat_MPIBAIJ*)A->data;
    167:   Mat            B = baij->B,Bnew;
    168:   Mat_SeqBAIJ    *Bbaij = (Mat_SeqBAIJ*)B->data;
    170:   PetscInt       i,j,mbs=Bbaij->mbs,n = A->N,col,*garray=baij->garray;
    171:   PetscInt       bs2 = baij->bs2,*nz,ec,m = A->m;
    172:   MatScalar      *a = Bbaij->a;
    173:   PetscScalar    *atmp;
    174: #if defined(PETSC_USE_MAT_SINGLE)
    175:   PetscInt       k;
    176: #endif

    179:   /* free stuff related to matrix-vec multiply */
    180:   VecGetSize(baij->lvec,&ec); /* needed for PetscLogObjectMemory below */
    181:   VecDestroy(baij->lvec); baij->lvec = 0;
    182:   VecScatterDestroy(baij->Mvctx); baij->Mvctx = 0;
    183:   if (baij->colmap) {
    184: #if defined (PETSC_USE_CTABLE)
    185:     PetscTableDelete(baij->colmap); baij->colmap = 0;
    186: #else
    187:     PetscFree(baij->colmap);
    188:     baij->colmap = 0;
    189:     PetscLogObjectMemory(A,-Bbaij->nbs*sizeof(PetscInt));
    190: #endif
    191:   }

    193:   /* make sure that B is assembled so we can access its values */
    194:   MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);
    195:   MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);

    197:   /* invent new B and copy stuff over */
    198:   PetscMalloc(mbs*sizeof(PetscInt),&nz);
    199:   for (i=0; i<mbs; i++) {
    200:     nz[i] = Bbaij->i[i+1]-Bbaij->i[i];
    201:   }
    202:   MatCreate(B->comm,&Bnew);
    203:   MatSetSizes(Bnew,m,n,m,n);
    204:   MatSetType(Bnew,B->type_name);
    205:   MatSeqBAIJSetPreallocation(Bnew,B->bs,0,nz);
    206:   MatSetOption(Bnew,MAT_COLUMN_ORIENTED);

    208: #if defined(PETSC_USE_MAT_SINGLE)
    209:   PetscMalloc(bs2*sizeof(PetscScalar),&atmp);
    210: #endif
    211:     for (i=0; i<mbs; i++) {
    212:       for (j=Bbaij->i[i]; j<Bbaij->i[i+1]; j++) {
    213:         col  = garray[Bbaij->j[j]];
    214: #if defined(PETSC_USE_MAT_SINGLE)
    215:         for (k=0; k<bs2; k++) atmp[k] = a[j*bs2+k];
    216: #else
    217:         atmp = a + j*bs2;
    218: #endif
    219:         MatSetValuesBlocked_SeqBAIJ(Bnew,1,&i,1,&col,atmp,B->insertmode);
    220:       }
    221:     }
    222:   MatSetOption(Bnew,MAT_ROW_ORIENTED);

    224: #if defined(PETSC_USE_MAT_SINGLE)
    225:   PetscFree(atmp);
    226: #endif

    228:   PetscFree(nz);
    229:   PetscFree(baij->garray);
    230:   baij->garray = 0;
    231:   PetscLogObjectMemory(A,-ec*sizeof(PetscInt));
    232:   MatDestroy(B);
    233:   PetscLogObjectParent(A,Bnew);
    234:   baij->B = Bnew;
    235:   A->was_assembled = PETSC_FALSE;
    236:   return(0);
    237: }

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

    241: static PetscInt *uglyrmapd = 0,*uglyrmapo = 0;  /* mapping from the local ordering to the "diagonal" and "off-diagonal"
    242:                                       parts of the local matrix */
    243: static Vec uglydd = 0,uglyoo = 0;   /* work vectors used to scale the two parts of the local matrix */

    248: PetscErrorCode MatMPIBAIJDiagonalScaleLocalSetUp(Mat inA,Vec scale)
    249: {
    250:   Mat_MPIBAIJ    *ina = (Mat_MPIBAIJ*) inA->data; /*access private part of matrix */
    251:   Mat_SeqBAIJ    *B = (Mat_SeqBAIJ*)ina->B->data;
    253:   PetscInt       bs = inA->bs,i,n,nt,j,cstart,cend,no,*garray = ina->garray,*lindices;
    254:   PetscInt       *r_rmapd,*r_rmapo;
    255:
    257:   MatGetOwnershipRange(inA,&cstart,&cend);
    258:   MatGetSize(ina->A,PETSC_NULL,&n);
    259:   PetscMalloc((inA->bmapping->n+1)*sizeof(PetscInt),&r_rmapd);
    260:   PetscMemzero(r_rmapd,inA->bmapping->n*sizeof(PetscInt));
    261:   nt   = 0;
    262:   for (i=0; i<inA->bmapping->n; i++) {
    263:     if (inA->bmapping->indices[i]*bs >= cstart && inA->bmapping->indices[i]*bs < cend) {
    264:       nt++;
    265:       r_rmapd[i] = inA->bmapping->indices[i] + 1;
    266:     }
    267:   }
    268:   if (nt*bs != n) SETERRQ2(PETSC_ERR_PLIB,"Hmm nt*bs %D n %D",nt*bs,n);
    269:   PetscMalloc((n+1)*sizeof(PetscInt),&uglyrmapd);
    270:   for (i=0; i<inA->bmapping->n; i++) {
    271:     if (r_rmapd[i]){
    272:       for (j=0; j<bs; j++) {
    273:         uglyrmapd[(r_rmapd[i]-1)*bs+j-cstart] = i*bs + j;
    274:       }
    275:     }
    276:   }
    277:   PetscFree(r_rmapd);
    278:   VecCreateSeq(PETSC_COMM_SELF,n,&uglydd);

    280:   PetscMalloc((ina->Nbs+1)*sizeof(PetscInt),&lindices);
    281:   PetscMemzero(lindices,ina->Nbs*sizeof(PetscInt));
    282:   for (i=0; i<B->nbs; i++) {
    283:     lindices[garray[i]] = i+1;
    284:   }
    285:   no   = inA->bmapping->n - nt;
    286:   PetscMalloc((inA->bmapping->n+1)*sizeof(PetscInt),&r_rmapo);
    287:   PetscMemzero(r_rmapo,inA->bmapping->n*sizeof(PetscInt));
    288:   nt   = 0;
    289:   for (i=0; i<inA->bmapping->n; i++) {
    290:     if (lindices[inA->bmapping->indices[i]]) {
    291:       nt++;
    292:       r_rmapo[i] = lindices[inA->bmapping->indices[i]];
    293:     }
    294:   }
    295:   if (nt > no) SETERRQ2(PETSC_ERR_PLIB,"Hmm nt %D no %D",nt,n);
    296:   PetscFree(lindices);
    297:   PetscMalloc((nt*bs+1)*sizeof(PetscInt),&uglyrmapo);
    298:   for (i=0; i<inA->bmapping->n; i++) {
    299:     if (r_rmapo[i]){
    300:       for (j=0; j<bs; j++) {
    301:         uglyrmapo[(r_rmapo[i]-1)*bs+j] = i*bs + j;
    302:       }
    303:     }
    304:   }
    305:   PetscFree(r_rmapo);
    306:   VecCreateSeq(PETSC_COMM_SELF,nt*bs,&uglyoo);

    308:   return(0);
    309: }

    313: PetscErrorCode PETSCMAT_DLLEXPORT MatMPIBAIJDiagonalScaleLocal(Mat A,Vec scale)
    314: {
    315:   /* This routine should really be abandoned as it duplicates MatDiagonalScaleLocal */
    316:   PetscErrorCode ierr,(*f)(Mat,Vec);

    319:   PetscObjectQueryFunction((PetscObject)A,"MatDiagonalScaleLocal_C",(void (**)(void))&f);
    320:   if (f) {
    321:     (*f)(A,scale);
    322:   }
    323:   return(0);
    324: }

    329: PetscErrorCode PETSCMAT_DLLEXPORT MatDiagonalScaleLocal_MPIBAIJ(Mat A,Vec scale)
    330: {
    331:   Mat_MPIBAIJ    *a = (Mat_MPIBAIJ*) A->data; /*access private part of matrix */
    333:   PetscInt       n,i;
    334:   PetscScalar    *d,*o,*s;
    335:
    337:   if (!uglyrmapd) {
    338:     MatMPIBAIJDiagonalScaleLocalSetUp(A,scale);
    339:   }

    341:   VecGetArray(scale,&s);
    342:
    343:   VecGetLocalSize(uglydd,&n);
    344:   VecGetArray(uglydd,&d);
    345:   for (i=0; i<n; i++) {
    346:     d[i] = s[uglyrmapd[i]]; /* copy "diagonal" (true local) portion of scale into dd vector */
    347:   }
    348:   VecRestoreArray(uglydd,&d);
    349:   /* column scale "diagonal" portion of local matrix */
    350:   MatDiagonalScale(a->A,PETSC_NULL,uglydd);

    352:   VecGetLocalSize(uglyoo,&n);
    353:   VecGetArray(uglyoo,&o);
    354:   for (i=0; i<n; i++) {
    355:     o[i] = s[uglyrmapo[i]]; /* copy "off-diagonal" portion of scale into oo vector */
    356:   }
    357:   VecRestoreArray(scale,&s);
    358:   VecRestoreArray(uglyoo,&o);
    359:   /* column scale "off-diagonal" portion of local matrix */
    360:   MatDiagonalScale(a->B,PETSC_NULL,uglyoo);

    362:   return(0);
    363: }