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 = PetscArraycpy(idxm,idx,m);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 = PetscArraycpy(idxm,idx,m);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 MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly 177 178 Not Collective 179 180 Input Arguments: 181 + A - Full matrix, generally parallel 182 . isrow - Local index set for the rows 183 - iscol - Local index set for the columns 184 185 Output Arguments: 186 . newmat - New serial Mat 187 188 Level: developer 189 190 Notes: 191 Most will use MatGetLocalSubMatrix() which returns a real matrix corresponding to the local 192 block if it available, such as with matrix formats that store these blocks separately. 193 194 The new matrix forwards MatSetValuesLocal() and MatSetValuesBlockedLocal() to the global system. 195 In general, it does not define MatMult() or any other functions. Local submatrices can be nested. 196 197 .seealso: MatSetValuesLocal(), MatSetValuesBlockedLocal(), MatGetLocalSubMatrix(), MatCreateSubMatrix() 198 @*/ 199 PetscErrorCode MatCreateLocalRef(Mat A,IS isrow,IS iscol,Mat *newmat) 200 { 201 PetscErrorCode ierr; 202 Mat_LocalRef *lr; 203 Mat B; 204 PetscInt m,n; 205 PetscBool islr; 206 207 PetscFunctionBegin; 208 PetscValidHeaderSpecific(A,MAT_CLASSID,1); 209 PetscValidHeaderSpecific(isrow,IS_CLASSID,2); 210 PetscValidHeaderSpecific(iscol,IS_CLASSID,3); 211 PetscValidPointer(newmat,4); 212 if (!A->rmap->mapping) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONGSTATE,"Matrix must have local to global mapping provided before this call"); 213 *newmat = NULL; 214 215 ierr = MatCreate(PETSC_COMM_SELF,&B);CHKERRQ(ierr); 216 ierr = ISGetLocalSize(isrow,&m);CHKERRQ(ierr); 217 ierr = ISGetLocalSize(iscol,&n);CHKERRQ(ierr); 218 ierr = MatSetSizes(B,m,n,m,n);CHKERRQ(ierr); 219 ierr = PetscObjectChangeTypeName((PetscObject)B,MATLOCALREF);CHKERRQ(ierr); 220 ierr = MatSetUp(B);CHKERRQ(ierr); 221 222 B->ops->destroy = MatDestroy_LocalRef; 223 224 ierr = PetscNewLog(B,&lr);CHKERRQ(ierr); 225 B->data = (void*)lr; 226 227 ierr = PetscObjectTypeCompare((PetscObject)A,MATLOCALREF,&islr);CHKERRQ(ierr); 228 if (islr) { 229 Mat_LocalRef *alr = (Mat_LocalRef*)A->data; 230 lr->Top = alr->Top; 231 } else { 232 /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */ 233 lr->Top = A; 234 } 235 { 236 ISLocalToGlobalMapping rltog,cltog; 237 PetscInt arbs,acbs,rbs,cbs; 238 239 /* We will translate directly to global indices for the top level */ 240 lr->SetValues = MatSetValues; 241 lr->SetValuesBlocked = MatSetValuesBlocked; 242 243 B->ops->setvalueslocal = MatSetValuesLocal_LocalRef_Scalar; 244 245 ierr = ISL2GCompose(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr); 246 if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) { 247 ierr = PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr); 248 cltog = rltog; 249 } else { 250 ierr = ISL2GCompose(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr); 251 } 252 /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK. Will be used later in 253 * MatSetValuesLocal_LocalRef_Scalar. */ 254 ierr = PetscObjectTypeCompare((PetscObject)isrow,ISBLOCK,&lr->rowisblock);CHKERRQ(ierr); 255 ierr = PetscObjectTypeCompare((PetscObject)iscol,ISBLOCK,&lr->colisblock);CHKERRQ(ierr); 256 ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr); 257 ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr); 258 ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr); 259 260 ierr = MatGetBlockSizes(A,&arbs,&acbs);CHKERRQ(ierr); 261 ierr = ISGetBlockSize(isrow,&rbs);CHKERRQ(ierr); 262 ierr = ISGetBlockSize(iscol,&cbs);CHKERRQ(ierr); 263 /* Always support block interface insertion on submatrix */ 264 ierr = PetscLayoutSetBlockSize(B->rmap,rbs);CHKERRQ(ierr); 265 ierr = PetscLayoutSetBlockSize(B->cmap,cbs);CHKERRQ(ierr); 266 if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) { 267 /* Top-level matrix has different block size, so we have to call its scalar insertion interface */ 268 B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar; 269 } else { 270 /* Block sizes match so we can forward values to the top level using the block interface */ 271 B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block; 272 273 ierr = ISL2GComposeBlock(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr); 274 if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) { 275 ierr = PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr); 276 cltog = rltog; 277 } else { 278 ierr = ISL2GComposeBlock(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr); 279 } 280 ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr); 281 ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr); 282 ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr); 283 } 284 } 285 *newmat = B; 286 PetscFunctionReturn(0); 287 } 288