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 PetscErrorCode MatMPIAIJDiagonalScaleLocalSetUp(Mat inA,Vec scale) 215 { 216 Mat_MPIAIJ *ina = (Mat_MPIAIJ*) inA->data; /*access private part of matrix */ 217 PetscErrorCode ierr; 218 PetscInt i,n,nt,cstart,cend,no,*garray = ina->garray,*lindices; 219 PetscInt *r_rmapd,*r_rmapo; 220 221 PetscFunctionBegin; 222 ierr = MatGetOwnershipRange(inA,&cstart,&cend);CHKERRQ(ierr); 223 ierr = MatGetSize(ina->A,NULL,&n);CHKERRQ(ierr); 224 ierr = PetscCalloc1(inA->rmap->mapping->n+1,&r_rmapd);CHKERRQ(ierr); 225 nt = 0; 226 for (i=0; i<inA->rmap->mapping->n; i++) { 227 if (inA->rmap->mapping->indices[i] >= cstart && inA->rmap->mapping->indices[i] < cend) { 228 nt++; 229 r_rmapd[i] = inA->rmap->mapping->indices[i] + 1; 230 } 231 } 232 if (nt != n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Hmm nt %D n %D",nt,n); 233 ierr = PetscMalloc1(n+1,&auglyrmapd);CHKERRQ(ierr); 234 for (i=0; i<inA->rmap->mapping->n; i++) { 235 if (r_rmapd[i]) { 236 auglyrmapd[(r_rmapd[i]-1)-cstart] = i; 237 } 238 } 239 ierr = PetscFree(r_rmapd);CHKERRQ(ierr); 240 ierr = VecCreateSeq(PETSC_COMM_SELF,n,&auglydd);CHKERRQ(ierr); 241 242 ierr = PetscCalloc1(inA->cmap->N+1,&lindices);CHKERRQ(ierr); 243 for (i=0; i<ina->B->cmap->n; i++) { 244 lindices[garray[i]] = i+1; 245 } 246 no = inA->rmap->mapping->n - nt; 247 ierr = PetscCalloc1(inA->rmap->mapping->n+1,&r_rmapo);CHKERRQ(ierr); 248 nt = 0; 249 for (i=0; i<inA->rmap->mapping->n; i++) { 250 if (lindices[inA->rmap->mapping->indices[i]]) { 251 nt++; 252 r_rmapo[i] = lindices[inA->rmap->mapping->indices[i]]; 253 } 254 } 255 if (nt > no) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Hmm nt %D no %D",nt,n); 256 ierr = PetscFree(lindices);CHKERRQ(ierr); 257 ierr = PetscMalloc1(nt+1,&auglyrmapo);CHKERRQ(ierr); 258 for (i=0; i<inA->rmap->mapping->n; i++) { 259 if (r_rmapo[i]) { 260 auglyrmapo[(r_rmapo[i]-1)] = i; 261 } 262 } 263 ierr = PetscFree(r_rmapo);CHKERRQ(ierr); 264 ierr = VecCreateSeq(PETSC_COMM_SELF,nt,&auglyoo);CHKERRQ(ierr); 265 PetscFunctionReturn(0); 266 } 267 268 PetscErrorCode MatMPIAIJDiagonalScaleLocal(Mat A,Vec scale) 269 { 270 /* This routine should really be abandoned as it duplicates MatDiagonalScaleLocal */ 271 PetscErrorCode ierr; 272 273 PetscFunctionBegin; 274 ierr = PetscTryMethod(A,"MatDiagonalScaleLocal_C",(Mat,Vec),(A,scale));CHKERRQ(ierr); 275 PetscFunctionReturn(0); 276 } 277 278 PetscErrorCode MatDiagonalScaleLocal_MPIAIJ(Mat A,Vec scale) 279 { 280 Mat_MPIAIJ *a = (Mat_MPIAIJ*) A->data; /*access private part of matrix */ 281 PetscErrorCode ierr; 282 PetscInt n,i; 283 PetscScalar *d,*o; 284 const PetscScalar *s; 285 286 PetscFunctionBegin; 287 if (!auglyrmapd) { 288 ierr = MatMPIAIJDiagonalScaleLocalSetUp(A,scale);CHKERRQ(ierr); 289 } 290 291 ierr = VecGetArrayRead(scale,&s);CHKERRQ(ierr); 292 293 ierr = VecGetLocalSize(auglydd,&n);CHKERRQ(ierr); 294 ierr = VecGetArray(auglydd,&d);CHKERRQ(ierr); 295 for (i=0; i<n; i++) { 296 d[i] = s[auglyrmapd[i]]; /* copy "diagonal" (true local) portion of scale into dd vector */ 297 } 298 ierr = VecRestoreArray(auglydd,&d);CHKERRQ(ierr); 299 /* column scale "diagonal" portion of local matrix */ 300 ierr = MatDiagonalScale(a->A,NULL,auglydd);CHKERRQ(ierr); 301 302 ierr = VecGetLocalSize(auglyoo,&n);CHKERRQ(ierr); 303 ierr = VecGetArray(auglyoo,&o);CHKERRQ(ierr); 304 for (i=0; i<n; i++) { 305 o[i] = s[auglyrmapo[i]]; /* copy "off-diagonal" portion of scale into oo vector */ 306 } 307 ierr = VecRestoreArrayRead(scale,&s);CHKERRQ(ierr); 308 ierr = VecRestoreArray(auglyoo,&o);CHKERRQ(ierr); 309 /* column scale "off-diagonal" portion of local matrix */ 310 ierr = MatDiagonalScale(a->B,NULL,auglyoo);CHKERRQ(ierr); 311 PetscFunctionReturn(0); 312 } 313