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 = PetscLogObjectParent((PetscObject)mat,(PetscObject)aij->Mvctx);CHKERRQ(ierr); 123 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)aij->lvec);CHKERRQ(ierr); 124 ierr = PetscLogObjectMemory((PetscObject)mat,ec*sizeof(PetscInt));CHKERRQ(ierr); 125 aij->garray = garray; 126 127 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)from);CHKERRQ(ierr); 128 ierr = PetscLogObjectParent((PetscObject)mat,(PetscObject)to);CHKERRQ(ierr); 129 130 ierr = ISDestroy(&from);CHKERRQ(ierr); 131 ierr = ISDestroy(&to);CHKERRQ(ierr); 132 ierr = VecDestroy(&gvec);CHKERRQ(ierr); 133 PetscFunctionReturn(0); 134 } 135 136 /* 137 Takes the local part of an already assembled MPIAIJ matrix 138 and disassembles it. This is to allow new nonzeros into the matrix 139 that require more communication in the matrix vector multiply. 140 Thus certain data-structures must be rebuilt. 141 142 Kind of slow! But that's what application programmers get when 143 they are sloppy. 144 */ 145 PetscErrorCode MatDisAssemble_MPIAIJ(Mat A) 146 { 147 Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 148 Mat B = aij->B,Bnew; 149 Mat_SeqAIJ *Baij = (Mat_SeqAIJ*)B->data; 150 PetscErrorCode ierr; 151 PetscInt i,j,m = B->rmap->n,n = A->cmap->N,col,ct = 0,*garray = aij->garray,*nz,ec; 152 PetscScalar v; 153 154 PetscFunctionBegin; 155 /* free stuff related to matrix-vec multiply */ 156 ierr = VecGetSize(aij->lvec,&ec);CHKERRQ(ierr); /* needed for PetscLogObjectMemory below */ 157 ierr = VecDestroy(&aij->lvec);CHKERRQ(ierr); 158 if (aij->colmap) { 159 #if defined(PETSC_USE_CTABLE) 160 ierr = PetscTableDestroy(&aij->colmap);CHKERRQ(ierr); 161 #else 162 ierr = PetscFree(aij->colmap);CHKERRQ(ierr); 163 ierr = PetscLogObjectMemory((PetscObject)A,-aij->B->cmap->n*sizeof(PetscInt));CHKERRQ(ierr); 164 #endif 165 } 166 167 /* make sure that B is assembled so we can access its values */ 168 ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 169 ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 170 171 /* invent new B and copy stuff over */ 172 ierr = PetscMalloc1(m+1,&nz);CHKERRQ(ierr); 173 for (i=0; i<m; i++) { 174 nz[i] = Baij->i[i+1] - Baij->i[i]; 175 } 176 ierr = MatCreate(PETSC_COMM_SELF,&Bnew);CHKERRQ(ierr); 177 ierr = MatSetSizes(Bnew,m,n,m,n);CHKERRQ(ierr); 178 ierr = MatSetBlockSizesFromMats(Bnew,A,A);CHKERRQ(ierr); 179 ierr = MatSetType(Bnew,((PetscObject)B)->type_name);CHKERRQ(ierr); 180 ierr = MatSeqAIJSetPreallocation(Bnew,0,nz);CHKERRQ(ierr); 181 182 if (Baij->nonew >= 0) { /* Inherit insertion error options (if positive). */ 183 ((Mat_SeqAIJ*)Bnew->data)->nonew = Baij->nonew; 184 } 185 186 /* 187 Ensure that B's nonzerostate is monotonically increasing. 188 Or should this follow the MatSetValues() loop to preserve B's nonzerstate across a MatDisAssemble() call? 189 */ 190 Bnew->nonzerostate = B->nonzerostate; 191 192 ierr = PetscFree(nz);CHKERRQ(ierr); 193 for (i=0; i<m; i++) { 194 for (j=Baij->i[i]; j<Baij->i[i+1]; j++) { 195 col = garray[Baij->j[ct]]; 196 v = Baij->a[ct++]; 197 ierr = MatSetValues(Bnew,1,&i,1,&col,&v,B->insertmode);CHKERRQ(ierr); 198 } 199 } 200 ierr = PetscFree(aij->garray);CHKERRQ(ierr); 201 ierr = PetscLogObjectMemory((PetscObject)A,-ec*sizeof(PetscInt));CHKERRQ(ierr); 202 ierr = MatDestroy(&B);CHKERRQ(ierr); 203 ierr = PetscLogObjectParent((PetscObject)A,(PetscObject)Bnew);CHKERRQ(ierr); 204 205 aij->B = Bnew; 206 A->was_assembled = PETSC_FALSE; 207 PetscFunctionReturn(0); 208 } 209 210 /* ugly stuff added for Glenn someday we should fix this up */ 211 212 static PetscInt *auglyrmapd = NULL,*auglyrmapo = NULL; /* mapping from the local ordering to the "diagonal" and "off-diagonal" parts of the local matrix */ 213 static Vec auglydd = NULL,auglyoo = NULL; /* work vectors used to scale the two parts of the local matrix */ 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