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