1 2 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/ 3 4 typedef struct { 5 Mat Top; 6 PetscBool rowisblock; 7 PetscBool colisblock; 8 PetscErrorCode (*SetValues)(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode); 9 PetscErrorCode (*SetValuesBlocked)(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode); 10 } Mat_LocalRef; 11 12 /* These need to be macros because they use sizeof */ 13 #define IndexSpaceGet(buf,nrow,ncol,irowm,icolm) do { \ 14 if (nrow + ncol > (PetscInt)(sizeof(buf)/sizeof(buf[0]))) { \ 15 ierr = PetscMalloc2(nrow,&irowm,ncol,&icolm);CHKERRQ(ierr); \ 16 } else { \ 17 irowm = &buf[0]; \ 18 icolm = &buf[nrow]; \ 19 } \ 20 } while (0) 21 22 #define IndexSpaceRestore(buf,nrow,ncol,irowm,icolm) do { \ 23 if (nrow + ncol > (PetscInt)(sizeof(buf)/sizeof(buf[0]))) { \ 24 ierr = PetscFree2(irowm,icolm);CHKERRQ(ierr); \ 25 } \ 26 } while (0) 27 28 static void BlockIndicesExpand(PetscInt n,const PetscInt idx[],PetscInt bs,PetscInt idxm[]) 29 { 30 PetscInt i,j; 31 for (i=0; i<n; i++) { 32 for (j=0; j<bs; j++) { 33 idxm[i*bs+j] = idx[i]*bs + j; 34 } 35 } 36 } 37 38 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Block(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv) 39 { 40 Mat_LocalRef *lr = (Mat_LocalRef*)A->data; 41 PetscErrorCode ierr; 42 PetscInt buf[4096],*irowm=NULL,*icolm; /* suppress maybe-uninitialized warning */ 43 44 PetscFunctionBegin; 45 if (!nrow || !ncol) PetscFunctionReturn(0); 46 IndexSpaceGet(buf,nrow,ncol,irowm,icolm); 47 ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr); 48 ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr); 49 ierr = (*lr->SetValuesBlocked)(lr->Top,nrow,irowm,ncol,icolm,y,addv);CHKERRQ(ierr); 50 IndexSpaceRestore(buf,nrow,ncol,irowm,icolm); 51 PetscFunctionReturn(0); 52 } 53 54 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv) 55 { 56 Mat_LocalRef *lr = (Mat_LocalRef*)A->data; 57 PetscErrorCode ierr; 58 PetscInt rbs,cbs,buf[4096],*irowm,*icolm; 59 60 PetscFunctionBegin; 61 ierr = MatGetBlockSizes(A,&rbs,&cbs);CHKERRQ(ierr); 62 IndexSpaceGet(buf,nrow*rbs,ncol*cbs,irowm,icolm); 63 BlockIndicesExpand(nrow,irow,rbs,irowm); 64 BlockIndicesExpand(ncol,icol,cbs,icolm); 65 ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow*rbs,irowm,irowm);CHKERRQ(ierr); 66 ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol*cbs,icolm,icolm);CHKERRQ(ierr); 67 ierr = (*lr->SetValues)(lr->Top,nrow*rbs,irowm,ncol*cbs,icolm,y,addv);CHKERRQ(ierr); 68 IndexSpaceRestore(buf,nrow*rbs,ncol*cbs,irowm,icolm); 69 PetscFunctionReturn(0); 70 } 71 72 static PetscErrorCode MatSetValuesLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv) 73 { 74 Mat_LocalRef *lr = (Mat_LocalRef*)A->data; 75 PetscErrorCode ierr; 76 PetscInt buf[4096],*irowm,*icolm; 77 78 PetscFunctionBegin; 79 IndexSpaceGet(buf,nrow,ncol,irowm,icolm); 80 /* If the row IS defining this submatrix was an ISBLOCK, then the unblocked LGMapApply is the right one to use. If 81 * instead it was (say) an ISSTRIDE with a block size > 1, then we need to use LGMapApplyBlock */ 82 if (lr->rowisblock) { 83 ierr = ISLocalToGlobalMappingApply(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr); 84 } else { 85 ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr); 86 } 87 /* As above, but for the column IS. */ 88 if (lr->colisblock) { 89 ierr = ISLocalToGlobalMappingApply(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr); 90 } else { 91 ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr); 92 } 93 ierr = (*lr->SetValues)(lr->Top,nrow,irowm,ncol,icolm,y,addv);CHKERRQ(ierr); 94 IndexSpaceRestore(buf,nrow,ncol,irowm,icolm); 95 PetscFunctionReturn(0); 96 } 97 98 /* Compose an IS with an ISLocalToGlobalMapping to map from IS source indices to global indices */ 99 static PetscErrorCode ISL2GCompose(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping *cltog) 100 { 101 PetscErrorCode ierr; 102 const PetscInt *idx; 103 PetscInt m,*idxm; 104 PetscInt bs; 105 PetscBool isblock; 106 107 PetscFunctionBegin; 108 PetscValidHeaderSpecific(is,IS_CLASSID,1); 109 PetscValidHeaderSpecific(ltog,IS_LTOGM_CLASSID,2); 110 PetscValidPointer(cltog,3); 111 ierr = PetscObjectTypeCompare((PetscObject)is,ISBLOCK,&isblock);CHKERRQ(ierr); 112 if (isblock) { 113 PetscInt lbs; 114 115 ierr = ISGetBlockSize(is,&bs);CHKERRQ(ierr); 116 ierr = ISLocalToGlobalMappingGetBlockSize(ltog,&lbs);CHKERRQ(ierr); 117 if (bs == lbs) { 118 ierr = ISGetLocalSize(is,&m);CHKERRQ(ierr); 119 m = m/bs; 120 ierr = ISBlockGetIndices(is,&idx);CHKERRQ(ierr); 121 ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr); 122 ierr = ISLocalToGlobalMappingApplyBlock(ltog,m,idx,idxm);CHKERRQ(ierr); 123 ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr); 124 ierr = ISBlockRestoreIndices(is,&idx);CHKERRQ(ierr); 125 PetscFunctionReturn(0); 126 } 127 } 128 ierr = ISGetLocalSize(is,&m);CHKERRQ(ierr); 129 ierr = ISGetIndices(is,&idx);CHKERRQ(ierr); 130 ierr = ISGetBlockSize(is,&bs);CHKERRQ(ierr); 131 ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr); 132 if (ltog) { 133 ierr = ISLocalToGlobalMappingApply(ltog,m,idx,idxm);CHKERRQ(ierr); 134 } else { 135 ierr = PetscMemcpy(idxm,idx,m*sizeof(PetscInt));CHKERRQ(ierr); 136 } 137 ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr); 138 ierr = ISRestoreIndices(is,&idx);CHKERRQ(ierr); 139 PetscFunctionReturn(0); 140 } 141 142 static PetscErrorCode ISL2GComposeBlock(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping *cltog) 143 { 144 PetscErrorCode ierr; 145 const PetscInt *idx; 146 PetscInt m,*idxm,bs; 147 148 PetscFunctionBegin; 149 PetscValidHeaderSpecific(is,IS_CLASSID,1); 150 PetscValidHeaderSpecific(ltog,IS_LTOGM_CLASSID,2); 151 PetscValidPointer(cltog,3); 152 ierr = ISBlockGetLocalSize(is,&m);CHKERRQ(ierr); 153 ierr = ISBlockGetIndices(is,&idx);CHKERRQ(ierr); 154 ierr = ISLocalToGlobalMappingGetBlockSize(ltog,&bs);CHKERRQ(ierr); 155 ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr); 156 if (ltog) { 157 ierr = ISLocalToGlobalMappingApplyBlock(ltog,m,idx,idxm);CHKERRQ(ierr); 158 } else { 159 ierr = PetscMemcpy(idxm,idx,m*sizeof(PetscInt));CHKERRQ(ierr); 160 } 161 ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr); 162 ierr = ISBlockRestoreIndices(is,&idx);CHKERRQ(ierr); 163 PetscFunctionReturn(0); 164 } 165 166 static PetscErrorCode MatDestroy_LocalRef(Mat B) 167 { 168 PetscErrorCode ierr; 169 170 PetscFunctionBegin; 171 ierr = PetscFree(B->data);CHKERRQ(ierr); 172 PetscFunctionReturn(0); 173 } 174 175 176 /*@ 177 MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly 178 179 Not Collective 180 181 Input Arguments: 182 + A - Full matrix, generally parallel 183 . isrow - Local index set for the rows 184 - iscol - Local index set for the columns 185 186 Output Arguments: 187 . newmat - New serial Mat 188 189 Level: developer 190 191 Notes: 192 Most will use MatGetLocalSubMatrix() which returns a real matrix corresponding to the local 193 block if it available, such as with matrix formats that store these blocks separately. 194 195 The new matrix forwards MatSetValuesLocal() and MatSetValuesBlockedLocal() to the global system. 196 In general, it does not define MatMult() or any other functions. Local submatrices can be nested. 197 198 .seealso: MatSetValuesLocal(), MatSetValuesBlockedLocal(), MatGetLocalSubMatrix(), MatCreateSubMatrix() 199 @*/ 200 PetscErrorCode MatCreateLocalRef(Mat A,IS isrow,IS iscol,Mat *newmat) 201 { 202 PetscErrorCode ierr; 203 Mat_LocalRef *lr; 204 Mat B; 205 PetscInt m,n; 206 PetscBool islr; 207 208 PetscFunctionBegin; 209 PetscValidHeaderSpecific(A,MAT_CLASSID,1); 210 PetscValidHeaderSpecific(isrow,IS_CLASSID,2); 211 PetscValidHeaderSpecific(iscol,IS_CLASSID,3); 212 PetscValidPointer(newmat,4); 213 if (!A->rmap->mapping) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONGSTATE,"Matrix must have local to global mapping provided before this call"); 214 *newmat = 0; 215 216 ierr = MatCreate(PETSC_COMM_SELF,&B);CHKERRQ(ierr); 217 ierr = ISGetLocalSize(isrow,&m);CHKERRQ(ierr); 218 ierr = ISGetLocalSize(iscol,&n);CHKERRQ(ierr); 219 ierr = MatSetSizes(B,m,n,m,n);CHKERRQ(ierr); 220 ierr = PetscObjectChangeTypeName((PetscObject)B,MATLOCALREF);CHKERRQ(ierr); 221 ierr = MatSetUp(B);CHKERRQ(ierr); 222 223 B->ops->destroy = MatDestroy_LocalRef; 224 225 ierr = PetscNewLog(B,&lr);CHKERRQ(ierr); 226 B->data = (void*)lr; 227 228 ierr = PetscObjectTypeCompare((PetscObject)A,MATLOCALREF,&islr);CHKERRQ(ierr); 229 if (islr) { 230 Mat_LocalRef *alr = (Mat_LocalRef*)A->data; 231 lr->Top = alr->Top; 232 } else { 233 /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */ 234 lr->Top = A; 235 } 236 { 237 ISLocalToGlobalMapping rltog,cltog; 238 PetscInt arbs,acbs,rbs,cbs; 239 240 /* We will translate directly to global indices for the top level */ 241 lr->SetValues = MatSetValues; 242 lr->SetValuesBlocked = MatSetValuesBlocked; 243 244 B->ops->setvalueslocal = MatSetValuesLocal_LocalRef_Scalar; 245 246 ierr = ISL2GCompose(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr); 247 if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) { 248 ierr = PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr); 249 cltog = rltog; 250 } else { 251 ierr = ISL2GCompose(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr); 252 } 253 /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK. Will be used later in 254 * MatSetValuesLocal_LocalRef_Scalar. */ 255 ierr = PetscObjectTypeCompare((PetscObject)isrow,ISBLOCK,&lr->rowisblock);CHKERRQ(ierr); 256 ierr = PetscObjectTypeCompare((PetscObject)iscol,ISBLOCK,&lr->colisblock);CHKERRQ(ierr); 257 ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr); 258 ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr); 259 ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr); 260 261 ierr = MatGetBlockSizes(A,&arbs,&acbs);CHKERRQ(ierr); 262 ierr = ISGetBlockSize(isrow,&rbs);CHKERRQ(ierr); 263 ierr = ISGetBlockSize(iscol,&cbs);CHKERRQ(ierr); 264 /* Always support block interface insertion on submatrix */ 265 ierr = PetscLayoutSetBlockSize(B->rmap,rbs);CHKERRQ(ierr); 266 ierr = PetscLayoutSetBlockSize(B->cmap,cbs);CHKERRQ(ierr); 267 if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) { 268 /* Top-level matrix has different block size, so we have to call its scalar insertion interface */ 269 B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar; 270 } else { 271 /* Block sizes match so we can forward values to the top level using the block interface */ 272 B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block; 273 274 ierr = ISL2GComposeBlock(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr); 275 if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) { 276 ierr = PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr); 277 cltog = rltog; 278 } else { 279 ierr = ISL2GComposeBlock(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr); 280 } 281 ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr); 282 ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr); 283 ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr); 284 } 285 } 286 *newmat = B; 287 PetscFunctionReturn(0); 288 } 289