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