1 2 /* 3 Support for the parallel AIJ matrix vector multiply 4 */ 5 #include <../src/mat/impls/aij/mpi/mpiaij.h> 6 #include <petsc-private/isimpl.h> /* needed because accesses data structure of ISLocalToGlobalMapping directly */ 7 8 #undef __FUNCT__ 9 #define __FUNCT__ "MatSetUpMultiply_MPIAIJ" 10 PetscErrorCode MatSetUpMultiply_MPIAIJ(Mat mat) 11 { 12 Mat_MPIAIJ *aij = (Mat_MPIAIJ*)mat->data; 13 Mat_SeqAIJ *B = (Mat_SeqAIJ*)(aij->B->data); 14 PetscErrorCode ierr; 15 PetscInt i,j,*aj = B->j,ec = 0,*garray; 16 IS from,to; 17 Vec gvec; 18 #if defined(PETSC_USE_CTABLE) 19 PetscTable gid1_lid1; 20 PetscTablePosition tpos; 21 PetscInt gid,lid; 22 #else 23 PetscInt N = mat->cmap->N,*indices; 24 #endif 25 26 PetscFunctionBegin; 27 #if defined(PETSC_USE_CTABLE) 28 /* use a table */ 29 ierr = PetscTableCreate(aij->B->rmap->n,mat->cmap->N+1,&gid1_lid1);CHKERRQ(ierr); 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 ierr = PetscTableFind(gid1_lid1,gid1,&data);CHKERRQ(ierr); 34 if (!data) { 35 /* one based table */ 36 ierr = PetscTableAdd(gid1_lid1,gid1,++ec,INSERT_VALUES);CHKERRQ(ierr); 37 } 38 } 39 } 40 /* form array of columns we need */ 41 ierr = PetscMalloc1((ec+1),&garray);CHKERRQ(ierr); 42 ierr = PetscTableGetHeadPosition(gid1_lid1,&tpos);CHKERRQ(ierr); 43 while (tpos) { 44 ierr = PetscTableGetNext(gid1_lid1,&tpos,&gid,&lid);CHKERRQ(ierr); 45 gid--; 46 lid--; 47 garray[lid] = gid; 48 } 49 ierr = PetscSortInt(ec,garray);CHKERRQ(ierr); /* sort, and rebuild */ 50 ierr = PetscTableRemoveAll(gid1_lid1);CHKERRQ(ierr); 51 for (i=0; i<ec; i++) { 52 ierr = PetscTableAdd(gid1_lid1,garray[i]+1,i+1,INSERT_VALUES);CHKERRQ(ierr); 53 } 54 /* compact out the extra columns in B */ 55 for (i=0; i<aij->B->rmap->n; i++) { 56 for (j=0; j<B->ilen[i]; j++) { 57 PetscInt gid1 = aj[B->i[i] + j] + 1; 58 ierr = PetscTableFind(gid1_lid1,gid1,&lid);CHKERRQ(ierr); 59 lid--; 60 aj[B->i[i] + j] = lid; 61 } 62 } 63 aij->B->cmap->n = aij->B->cmap->N = ec; 64 65 ierr = PetscLayoutSetUp((aij->B->cmap));CHKERRQ(ierr); 66 ierr = PetscTableDestroy(&gid1_lid1);CHKERRQ(ierr); 67 #else 68 /* Make an array as long as the number of columns */ 69 /* mark those columns that are in aij->B */ 70 ierr = PetscCalloc1(N+1,&indices);CHKERRQ(ierr); 71 for (i=0; i<aij->B->rmap->n; i++) { 72 for (j=0; j<B->ilen[i]; j++) { 73 if (!indices[aj[B->i[i] + j]]) ec++; 74 indices[aj[B->i[i] + j]] = 1; 75 } 76 } 77 78 /* form array of columns we need */ 79 ierr = PetscMalloc1((ec+1),&garray);CHKERRQ(ierr); 80 ec = 0; 81 for (i=0; i<N; i++) { 82 if (indices[i]) garray[ec++] = i; 83 } 84 85 /* make indices now point into garray */ 86 for (i=0; i<ec; i++) { 87 indices[garray[i]] = i; 88 } 89 90 /* compact out the extra columns in B */ 91 for (i=0; i<aij->B->rmap->n; i++) { 92 for (j=0; j<B->ilen[i]; j++) { 93 aj[B->i[i] + j] = indices[aj[B->i[i] + j]]; 94 } 95 } 96 aij->B->cmap->n = aij->B->cmap->N = ec; 97 aij->B->cmap->bs = 1; 98 99 ierr = PetscLayoutSetUp((aij->B->cmap));CHKERRQ(ierr); 100 ierr = PetscFree(indices);CHKERRQ(ierr); 101 #endif 102 /* create local vector that is used to scatter into */ 103 ierr = VecCreateSeq(PETSC_COMM_SELF,ec,&aij->lvec);CHKERRQ(ierr); 104 105 /* create two temporary Index sets for build scatter gather */ 106 ierr = ISCreateGeneral(((PetscObject)mat)->comm,ec,garray,PETSC_COPY_VALUES,&from);CHKERRQ(ierr); 107 108 ierr = ISCreateStride(PETSC_COMM_SELF,ec,0,1,&to);CHKERRQ(ierr); 109 110 /* create temporary global vector to generate scatter context */ 111 /* This does not allocate the array's memory so is efficient */ 112 ierr = VecCreateMPIWithArray(PetscObjectComm((PetscObject)mat),1,mat->cmap->n,mat->cmap->N,NULL,&gvec);CHKERRQ(ierr); 113 114 /* generate the scatter context */ 115 ierr = VecScatterCreate(gvec,from,aij->lvec,to,&aij->Mvctx);CHKERRQ(ierr); 116 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)aij->Mvctx);CHKERRQ(ierr); 117 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)aij->lvec);CHKERRQ(ierr); 118 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)from);CHKERRQ(ierr); 119 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)to);CHKERRQ(ierr); 120 121 aij->garray = garray; 122 123 ierr = PetscLogObjectMemory((PetscObject)mat,(ec+1)*sizeof(PetscInt));CHKERRQ(ierr); 124 ierr = ISDestroy(&from);CHKERRQ(ierr); 125 ierr = ISDestroy(&to);CHKERRQ(ierr); 126 ierr = VecDestroy(&gvec);CHKERRQ(ierr); 127 PetscFunctionReturn(0); 128 } 129 130 131 #undef __FUNCT__ 132 #define __FUNCT__ "MatDisAssemble_MPIAIJ" 133 /* 134 Takes the local part of an already assembled MPIAIJ matrix 135 and disassembles it. This is to allow new nonzeros into the matrix 136 that require more communication in the matrix vector multiply. 137 Thus certain data-structures must be rebuilt. 138 139 Kind of slow! But that's what application programmers get when 140 they are sloppy. 141 */ 142 PetscErrorCode MatDisAssemble_MPIAIJ(Mat A) 143 { 144 Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 145 Mat B = aij->B,Bnew; 146 Mat_SeqAIJ *Baij = (Mat_SeqAIJ*)B->data; 147 PetscErrorCode ierr; 148 PetscInt i,j,m = B->rmap->n,n = A->cmap->N,col,ct = 0,*garray = aij->garray,*nz,ec; 149 PetscScalar v; 150 151 PetscFunctionBegin; 152 /* free stuff related to matrix-vec multiply */ 153 ierr = VecGetSize(aij->lvec,&ec);CHKERRQ(ierr); /* needed for PetscLogObjectMemory below */ 154 ierr = VecDestroy(&aij->lvec);CHKERRQ(ierr); 155 ierr = VecScatterDestroy(&aij->Mvctx);CHKERRQ(ierr); 156 if (aij->colmap) { 157 #if defined(PETSC_USE_CTABLE) 158 ierr = PetscTableDestroy(&aij->colmap);CHKERRQ(ierr); 159 #else 160 ierr = PetscFree(aij->colmap);CHKERRQ(ierr); 161 ierr = PetscLogObjectMemory((PetscObject)A,-aij->B->cmap->n*sizeof(PetscInt));CHKERRQ(ierr); 162 #endif 163 } 164 165 /* make sure that B is assembled so we can access its values */ 166 ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 167 ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 168 169 /* invent new B and copy stuff over */ 170 ierr = PetscMalloc1((m+1),&nz);CHKERRQ(ierr); 171 for (i=0; i<m; i++) { 172 nz[i] = Baij->i[i+1] - Baij->i[i]; 173 } 174 ierr = MatCreate(PETSC_COMM_SELF,&Bnew);CHKERRQ(ierr); 175 ierr = MatSetSizes(Bnew,m,n,m,n);CHKERRQ(ierr); 176 ierr = MatSetBlockSizes(Bnew,A->rmap->bs,A->cmap->bs);CHKERRQ(ierr); 177 ierr = MatSetType(Bnew,((PetscObject)B)->type_name);CHKERRQ(ierr); 178 ierr = MatSeqAIJSetPreallocation(Bnew,0,nz);CHKERRQ(ierr); 179 180 ((Mat_SeqAIJ*)Bnew->data)->nonew = Baij->nonew; /* Inherit insertion error options. */ 181 182 ierr = PetscFree(nz);CHKERRQ(ierr); 183 for (i=0; i<m; i++) { 184 for (j=Baij->i[i]; j<Baij->i[i+1]; j++) { 185 col = garray[Baij->j[ct]]; 186 v = Baij->a[ct++]; 187 ierr = MatSetValues(Bnew,1,&i,1,&col,&v,B->insertmode);CHKERRQ(ierr); 188 } 189 } 190 ierr = PetscFree(aij->garray);CHKERRQ(ierr); 191 ierr = PetscLogObjectMemory((PetscObject)A,-ec*sizeof(PetscInt));CHKERRQ(ierr); 192 ierr = MatDestroy(&B);CHKERRQ(ierr); 193 ierr = PetscLogObjectParent((PetscObject)A,(PetscObject)Bnew);CHKERRQ(ierr); 194 195 aij->B = Bnew; 196 A->was_assembled = PETSC_FALSE; 197 PetscFunctionReturn(0); 198 } 199 200 /* ugly stuff added for Glenn someday we should fix this up */ 201 202 static PetscInt *auglyrmapd = 0,*auglyrmapo = 0; /* mapping from the local ordering to the "diagonal" and "off-diagonal" parts of the local matrix */ 203 static Vec auglydd = 0,auglyoo = 0; /* work vectors used to scale the two parts of the local matrix */ 204 205 206 #undef __FUNCT__ 207 #define __FUNCT__ "MatMPIAIJDiagonalScaleLocalSetUp" 208 PetscErrorCode MatMPIAIJDiagonalScaleLocalSetUp(Mat inA,Vec scale) 209 { 210 Mat_MPIAIJ *ina = (Mat_MPIAIJ*) inA->data; /*access private part of matrix */ 211 PetscErrorCode ierr; 212 PetscInt i,n,nt,cstart,cend,no,*garray = ina->garray,*lindices; 213 PetscInt *r_rmapd,*r_rmapo; 214 215 PetscFunctionBegin; 216 ierr = MatGetOwnershipRange(inA,&cstart,&cend);CHKERRQ(ierr); 217 ierr = MatGetSize(ina->A,NULL,&n);CHKERRQ(ierr); 218 ierr = PetscCalloc1(inA->rmap->mapping->n+1,&r_rmapd);CHKERRQ(ierr); 219 nt = 0; 220 for (i=0; i<inA->rmap->mapping->n; i++) { 221 if (inA->rmap->mapping->indices[i] >= cstart && inA->rmap->mapping->indices[i] < cend) { 222 nt++; 223 r_rmapd[i] = inA->rmap->mapping->indices[i] + 1; 224 } 225 } 226 if (nt != n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Hmm nt %D n %D",nt,n); 227 ierr = PetscMalloc1((n+1),&auglyrmapd);CHKERRQ(ierr); 228 for (i=0; i<inA->rmap->mapping->n; i++) { 229 if (r_rmapd[i]) { 230 auglyrmapd[(r_rmapd[i]-1)-cstart] = i; 231 } 232 } 233 ierr = PetscFree(r_rmapd);CHKERRQ(ierr); 234 ierr = VecCreateSeq(PETSC_COMM_SELF,n,&auglydd);CHKERRQ(ierr); 235 236 ierr = PetscCalloc1(inA->cmap->N+1,&lindices);CHKERRQ(ierr); 237 for (i=0; i<ina->B->cmap->n; i++) { 238 lindices[garray[i]] = i+1; 239 } 240 no = inA->rmap->mapping->n - nt; 241 ierr = PetscCalloc1(inA->rmap->mapping->n+1,&r_rmapo);CHKERRQ(ierr); 242 nt = 0; 243 for (i=0; i<inA->rmap->mapping->n; i++) { 244 if (lindices[inA->rmap->mapping->indices[i]]) { 245 nt++; 246 r_rmapo[i] = lindices[inA->rmap->mapping->indices[i]]; 247 } 248 } 249 if (nt > no) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Hmm nt %D no %D",nt,n); 250 ierr = PetscFree(lindices);CHKERRQ(ierr); 251 ierr = PetscMalloc1((nt+1),&auglyrmapo);CHKERRQ(ierr); 252 for (i=0; i<inA->rmap->mapping->n; i++) { 253 if (r_rmapo[i]) { 254 auglyrmapo[(r_rmapo[i]-1)] = i; 255 } 256 } 257 ierr = PetscFree(r_rmapo);CHKERRQ(ierr); 258 ierr = VecCreateSeq(PETSC_COMM_SELF,nt,&auglyoo);CHKERRQ(ierr); 259 PetscFunctionReturn(0); 260 } 261 262 #undef __FUNCT__ 263 #define __FUNCT__ "MatMPIAIJDiagonalScaleLocal" 264 PetscErrorCode MatMPIAIJDiagonalScaleLocal(Mat A,Vec scale) 265 { 266 /* This routine should really be abandoned as it duplicates MatDiagonalScaleLocal */ 267 PetscErrorCode ierr; 268 269 PetscFunctionBegin; 270 ierr = PetscTryMethod(A,"MatDiagonalScaleLocal_C",(Mat,Vec),(A,scale));CHKERRQ(ierr); 271 PetscFunctionReturn(0); 272 } 273 274 #undef __FUNCT__ 275 #define __FUNCT__ "MatDiagonalScaleLocal_MPIAIJ" 276 PetscErrorCode MatDiagonalScaleLocal_MPIAIJ(Mat A,Vec scale) 277 { 278 Mat_MPIAIJ *a = (Mat_MPIAIJ*) A->data; /*access private part of matrix */ 279 PetscErrorCode ierr; 280 PetscInt n,i; 281 PetscScalar *d,*o,*s; 282 283 PetscFunctionBegin; 284 if (!auglyrmapd) { 285 ierr = MatMPIAIJDiagonalScaleLocalSetUp(A,scale);CHKERRQ(ierr); 286 } 287 288 ierr = VecGetArray(scale,&s);CHKERRQ(ierr); 289 290 ierr = VecGetLocalSize(auglydd,&n);CHKERRQ(ierr); 291 ierr = VecGetArray(auglydd,&d);CHKERRQ(ierr); 292 for (i=0; i<n; i++) { 293 d[i] = s[auglyrmapd[i]]; /* copy "diagonal" (true local) portion of scale into dd vector */ 294 } 295 ierr = VecRestoreArray(auglydd,&d);CHKERRQ(ierr); 296 /* column scale "diagonal" portion of local matrix */ 297 ierr = MatDiagonalScale(a->A,NULL,auglydd);CHKERRQ(ierr); 298 299 ierr = VecGetLocalSize(auglyoo,&n);CHKERRQ(ierr); 300 ierr = VecGetArray(auglyoo,&o);CHKERRQ(ierr); 301 for (i=0; i<n; i++) { 302 o[i] = s[auglyrmapo[i]]; /* copy "off-diagonal" portion of scale into oo vector */ 303 } 304 ierr = VecRestoreArray(scale,&s);CHKERRQ(ierr); 305 ierr = VecRestoreArray(auglyoo,&o);CHKERRQ(ierr); 306 /* column scale "off-diagonal" portion of local matrix */ 307 ierr = MatDiagonalScale(a->B,NULL,auglyoo);CHKERRQ(ierr); 308 PetscFunctionReturn(0); 309 } 310 311 312