Actual source code: mpb_aij.c
1: #include "../src/mat/impls/aij/mpi/mpiaij.h"
3: /*
5: This routine creates multiple [bjacobi] 'parallel submatrices' from
6: a given 'mat' object. Each submatrix can span multiple procs.
8: The submatrix partition across processors is dicated by 'subComm' a
9: communicator obtained by com_split(comm). Note: the comm_split
10: is not restriced to be grouped with consequitive original ranks.
12: Due the comm_split() usage, the parallel layout of the submatrices
13: map directly to the layout of the original matrix [wrt the local
14: row,col partitioning]. So the original 'DiagonalMat' naturally maps
15: into the 'DiagonalMat' of the subMat, hence it is used directly from
16: the subMat. However the offDiagMat looses some columns - and this is
17: reconstructed with MatSetValues()
18:
19: */
23: PetscErrorCode MatGetMultiProcBlock_MPIAIJ(Mat mat, MPI_Comm subComm, Mat* subMat)
24: {
26: Mat_MPIAIJ *aij = (Mat_MPIAIJ*)mat->data;
27: Mat_SeqAIJ* aijB = (Mat_SeqAIJ*)aij->B->data;
28: PetscMPIInt commRank,subCommSize,subCommRank;
29: PetscMPIInt *commRankMap,subRank,rank;
30: PetscInt *garrayCMap,col,i,j,*nnz,newRow,newCol;
34: /* create subMat object with the relavent layout */
35: MatCreate(subComm,subMat);
36: MatSetType(*subMat,MATMPIAIJ);
37: MatSetSizes(*subMat,mat->rmap->n,mat->cmap->n,PETSC_DECIDE,PETSC_DECIDE);
38: /* need to setup rmap and cmap before Preallocation */
39: PetscLayoutSetBlockSize((*subMat)->rmap,mat->rmap->bs);
40: PetscLayoutSetBlockSize((*subMat)->cmap,mat->cmap->bs);
41: PetscLayoutSetUp((*subMat)->rmap);
42: PetscLayoutSetUp((*subMat)->cmap);
44: /* create a map of comm_rank from subComm to comm */
45: MPI_Comm_rank(((PetscObject)mat)->comm,&commRank);
46: MPI_Comm_size(subComm,&subCommSize);
47: MPI_Comm_rank(subComm,&subCommRank);
48: PetscMalloc(subCommSize*sizeof(PetscMPIInt),&commRankMap);
49: MPI_Allgather(&commRank,1,MPI_INT,commRankMap,1,MPI_INT,subComm);
51: /* Traverse garray and identify column indices [of offdiag mat] that
52: should be discarded. For the ones not discarded, store the newCol+1
53: value in garrayCMap */
54: PetscMalloc(aij->B->cmap->n*sizeof(PetscInt),&garrayCMap);
55: PetscMemzero(garrayCMap,aij->B->cmap->n*sizeof(PetscInt));
56: for (i=0; i<aij->B->cmap->n; i++) {
57: col = aij->garray[i];
58: for (subRank=0; subRank<subCommSize; subRank++) {
59: rank = commRankMap[subRank];
60: if ((col >= mat->cmap->range[rank]) && (col < mat->cmap->range[rank+1])) {
61: garrayCMap[i] = (*subMat)->cmap->range[subRank] + col - mat->cmap->range[rank]+1;
62: break;
63: }
64: }
65: }
67: /* Now compute preallocation for the offdiag mat */
68: PetscMalloc(aij->B->rmap->n*sizeof(PetscInt),&nnz);
69: PetscMemzero(nnz,aij->B->rmap->n*sizeof(PetscInt));
70: for (i=0; i<aij->B->rmap->n; i++) {
71: for (j=aijB->i[i]; j<aijB->i[i+1]; j++) {
72: if (garrayCMap[aijB->j[j]]) nnz[i]++;
73: }
74: }
75: MatMPIAIJSetPreallocation(*(subMat),PETSC_NULL,PETSC_NULL,PETSC_NULL,nnz);
77: /* reuse diag block with the new submat */
78: MatDestroy(((Mat_MPIAIJ*)((*subMat)->data))->A);
79: ((Mat_MPIAIJ*)((*subMat)->data))->A = aij->A;
80: PetscObjectReference((PetscObject)aij->A);
82: /* Now traverse aij->B and insert values into subMat */
83: for (i=0; i<aij->B->rmap->n; i++) {
84: newRow = (*subMat)->rmap->range[subCommRank] + i;
85: for (j=aijB->i[i]; j<aijB->i[i+1]; j++) {
86: newCol = garrayCMap[aijB->j[j]];
87: if (newCol) {
88: newCol--; /* remove the increment */
89: MatSetValues(*subMat,1,&newRow,1,&newCol,(aijB->a+j),INSERT_VALUES);
90: }
91: }
92: }
94: /* assemble the submat */
95: MatAssemblyBegin(*subMat,MAT_FINAL_ASSEMBLY);
96: MatAssemblyEnd(*subMat,MAT_FINAL_ASSEMBLY);
97:
98: /* deallocate temporary data */
99: PetscFree(commRankMap);
100: PetscFree(garrayCMap);
101: PetscFree(nnz);
103: return(0);
104: }