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 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: `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(PetscNewLog(B, &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