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