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