1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/ 2 #include <petsc/private/hashsetij.h> 3 4 typedef struct { 5 PetscHSetIJ ht; 6 PetscInt *dnz, *onz; 7 PetscInt *dnzu, *onzu; 8 PetscBool nooffproc; 9 PetscBool used; 10 } Mat_Preallocator; 11 12 PetscErrorCode MatDestroy_Preallocator(Mat A) { 13 Mat_Preallocator *p = (Mat_Preallocator *)A->data; 14 15 PetscFunctionBegin; 16 PetscCall(MatStashDestroy_Private(&A->stash)); 17 PetscCall(PetscHSetIJDestroy(&p->ht)); 18 PetscCall(PetscFree4(p->dnz, p->onz, p->dnzu, p->onzu)); 19 PetscCall(PetscFree(A->data)); 20 PetscCall(PetscObjectChangeTypeName((PetscObject)A, NULL)); 21 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", NULL)); 22 PetscFunctionReturn(0); 23 } 24 25 PetscErrorCode MatSetUp_Preallocator(Mat A) { 26 Mat_Preallocator *p = (Mat_Preallocator *)A->data; 27 PetscInt m, bs, mbs; 28 29 PetscFunctionBegin; 30 PetscCall(PetscLayoutSetUp(A->rmap)); 31 PetscCall(PetscLayoutSetUp(A->cmap)); 32 PetscCall(MatGetLocalSize(A, &m, NULL)); 33 PetscCall(PetscHSetIJCreate(&p->ht)); 34 PetscCall(MatGetBlockSize(A, &bs)); 35 /* Do not bother bstash since MatPreallocator does not implement MatSetValuesBlocked */ 36 PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), 1, &A->stash)); 37 /* arrays are for blocked rows/cols */ 38 mbs = m / bs; 39 PetscCall(PetscCalloc4(mbs, &p->dnz, mbs, &p->onz, mbs, &p->dnzu, mbs, &p->onzu)); 40 PetscFunctionReturn(0); 41 } 42 43 PetscErrorCode MatSetValues_Preallocator(Mat A, PetscInt m, const PetscInt *rows, PetscInt n, const PetscInt *cols, const PetscScalar *values, InsertMode addv) { 44 Mat_Preallocator *p = (Mat_Preallocator *)A->data; 45 PetscInt rStart, rEnd, r, cStart, cEnd, c, bs; 46 47 PetscFunctionBegin; 48 PetscCall(MatGetBlockSize(A, &bs)); 49 PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd)); 50 PetscCall(MatGetOwnershipRangeColumn(A, &cStart, &cEnd)); 51 for (r = 0; r < m; ++r) { 52 PetscHashIJKey key; 53 PetscBool missing; 54 55 key.i = rows[r]; 56 if (key.i < 0) continue; 57 if ((key.i < rStart) || (key.i >= rEnd)) { 58 PetscCall(MatStashValuesRow_Private(&A->stash, key.i, n, cols, values, PETSC_FALSE)); 59 } else { /* Hash table is for blocked rows/cols */ 60 key.i = rows[r] / bs; 61 for (c = 0; c < n; ++c) { 62 key.j = cols[c] / bs; 63 if (key.j < 0) continue; 64 PetscCall(PetscHSetIJQueryAdd(p->ht, key, &missing)); 65 if (missing) { 66 if ((key.j >= cStart / bs) && (key.j < cEnd / bs)) { 67 ++p->dnz[key.i - rStart / bs]; 68 if (key.j >= key.i) ++p->dnzu[key.i - rStart / bs]; 69 } else { 70 ++p->onz[key.i - rStart / bs]; 71 if (key.j >= key.i) ++p->onzu[key.i - rStart / bs]; 72 } 73 } 74 } 75 } 76 } 77 PetscFunctionReturn(0); 78 } 79 80 PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type) { 81 PetscInt nstash, reallocs; 82 83 PetscFunctionBegin; 84 PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range)); 85 PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs)); 86 PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs)); 87 PetscFunctionReturn(0); 88 } 89 90 PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type) { 91 PetscScalar *val; 92 PetscInt *row, *col; 93 PetscInt i, j, rstart, ncols, flg; 94 PetscMPIInt n; 95 Mat_Preallocator *p = (Mat_Preallocator *)A->data; 96 97 PetscFunctionBegin; 98 p->nooffproc = PETSC_TRUE; 99 while (1) { 100 PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg)); 101 if (flg) p->nooffproc = PETSC_FALSE; 102 if (!flg) break; 103 104 for (i = 0; i < n;) { 105 /* Now identify the consecutive vals belonging to the same row */ 106 for (j = i, rstart = row[j]; j < n; j++) { 107 if (row[j] != rstart) break; 108 } 109 if (j < n) ncols = j - i; 110 else ncols = n - i; 111 /* Now assemble all these values with a single function call */ 112 PetscCall(MatSetValues_Preallocator(A, 1, row + i, ncols, col + i, val + i, INSERT_VALUES)); 113 i = j; 114 } 115 } 116 PetscCall(MatStashScatterEnd_Private(&A->stash)); 117 PetscCall(MPIU_Allreduce(MPI_IN_PLACE, &p->nooffproc, 1, MPIU_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A))); 118 PetscFunctionReturn(0); 119 } 120 121 PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer) { 122 PetscFunctionBegin; 123 PetscFunctionReturn(0); 124 } 125 126 PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg) { 127 PetscFunctionBegin; 128 PetscFunctionReturn(0); 129 } 130 131 PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A) { 132 Mat_Preallocator *p = (Mat_Preallocator *)mat->data; 133 PetscInt bs; 134 135 PetscFunctionBegin; 136 PetscCheck(!p->used, PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "MatPreallocatorPreallocate() can only be used once for a give MatPreallocator object. Consider using MatDuplicate() after preallocation."); 137 p->used = PETSC_TRUE; 138 if (!fill) PetscCall(PetscHSetIJDestroy(&p->ht)); 139 PetscCall(MatGetBlockSize(mat, &bs)); 140 PetscCall(MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu)); 141 PetscCall(MatSetUp(A)); 142 PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE)); 143 if (fill) { 144 PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc)); 145 PetscHashIter hi; 146 PetscHashIJKey key; 147 PetscScalar *zeros; 148 PetscInt n, maxrow = 1, *cols, rStart, rEnd, *rowstarts; 149 150 PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd)); 151 // Ownership range is in terms of scalar entries, but we deal with blocks 152 rStart /= bs; 153 rEnd /= bs; 154 PetscCall(PetscHSetIJGetSize(p->ht, &n)); 155 PetscCall(PetscMalloc2(n, &cols, rEnd - rStart + 1, &rowstarts)); 156 rowstarts[0] = 0; 157 for (PetscInt i = 0; i < rEnd - rStart; i++) { 158 rowstarts[i + 1] = rowstarts[i] + p->dnz[i] + p->onz[i]; 159 maxrow = PetscMax(maxrow, p->dnz[i] + p->onz[i]); 160 } 161 PetscCheck(rowstarts[rEnd - rStart] == n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Hash claims %" PetscInt_FMT " entries, but dnz+onz counts %" PetscInt_FMT, n, rowstarts[rEnd - rStart]); 162 163 PetscHashIterBegin(p->ht, hi); 164 while (!PetscHashIterAtEnd(p->ht, hi)) { 165 PetscHashIterGetKey(p->ht, hi, key); 166 PetscInt lrow = key.i - rStart; 167 cols[rowstarts[lrow]] = key.j; 168 rowstarts[lrow]++; 169 PetscHashIterNext(p->ht, hi); 170 } 171 PetscCall(PetscHSetIJDestroy(&p->ht)); 172 173 PetscCall(PetscCalloc1(maxrow * bs * bs, &zeros)); 174 for (PetscInt i = 0; i < rEnd - rStart; i++) { 175 PetscInt grow = rStart + i; 176 PetscInt end = rowstarts[i], start = end - p->dnz[i] - p->onz[i]; 177 PetscCall(PetscSortInt(end - start, &cols[start])); 178 PetscCall(MatSetValuesBlocked(A, 1, &grow, end - start, &cols[start], zeros, INSERT_VALUES)); 179 } 180 PetscCall(PetscFree(zeros)); 181 PetscCall(PetscFree2(cols, rowstarts)); 182 183 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 184 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 185 PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE)); 186 } 187 PetscFunctionReturn(0); 188 } 189 190 /*@ 191 MatPreallocatorPreallocate - Preallocates the A matrix, using information from a `MATPREALLOCATOR` mat, optionally filling A with zeros 192 193 Input Parameters: 194 + mat - the `MATPREALLOCATOR` preallocator matrix 195 . fill - fill the matrix with zeros 196 - A - the matrix to be preallocated 197 198 Notes: 199 This `MatType` implementation provides a helper utility to define the correct 200 preallocation data for a given nonzero structure. Use this object like a 201 regular matrix, e.g. loop over the nonzero structure of the matrix and 202 call `MatSetValues()` or `MatSetValuesBlocked()` to indicate the nonzero locations. 203 The matrix entries provided to `MatSetValues()` will be ignored, it only uses 204 the row / col indices provided to determine the information required to be 205 passed to `MatXAIJSetPreallocation()`. Once you have looped over the nonzero 206 structure, you must call `MatAssemblyBegin()`, `MatAssemblyEnd()` on mat. 207 208 After you have assembled the preallocator matrix (mat), call `MatPreallocatorPreallocate()` 209 to define the preallocation information on the matrix (A). Setting the parameter 210 fill = PETSC_TRUE will insert zeros into the matrix A. Internally `MatPreallocatorPreallocate()` 211 will call `MatSetOption`(A, `MAT_NEW_NONZERO_ALLOCATION_ERR`, `PETSC_TRUE`); 212 213 This function may only be called once for a given `MATPREALLOCATOR` object. If 214 multiple `Mat`s need to be preallocated, consider using `MatDuplicate()` after 215 this function. 216 217 Level: advanced 218 219 .seealso: `MATPREALLOCATOR`, `MatXAIJSetPreallocation()` 220 @*/ 221 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A) { 222 PetscFunctionBegin; 223 PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 224 PetscValidLogicalCollectiveBool(mat, fill, 2); 225 PetscValidHeaderSpecific(A, MAT_CLASSID, 3); 226 PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat, PetscBool, Mat), (mat, fill, A)); 227 PetscFunctionReturn(0); 228 } 229 230 /*MC 231 MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation. 232 233 Operations Provided: 234 .vb 235 MatSetValues() 236 .ve 237 238 Options Database Keys: 239 . -mat_type preallocator - sets the matrix type to `MATPREALLOCATOR` during a call to `MatSetFromOptions()` 240 241 Level: advanced 242 243 .seealso: `MATPREALLOCATOR`, `Mat`, `MatPreallocatorPreallocate()` 244 M*/ 245 246 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) { 247 Mat_Preallocator *p; 248 249 PetscFunctionBegin; 250 PetscCall(PetscNew(&p)); 251 A->data = (void *)p; 252 253 p->ht = NULL; 254 p->dnz = NULL; 255 p->onz = NULL; 256 p->dnzu = NULL; 257 p->onzu = NULL; 258 p->used = PETSC_FALSE; 259 260 /* matrix ops */ 261 PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps))); 262 263 A->ops->destroy = MatDestroy_Preallocator; 264 A->ops->setup = MatSetUp_Preallocator; 265 A->ops->setvalues = MatSetValues_Preallocator; 266 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 267 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 268 A->ops->view = MatView_Preallocator; 269 A->ops->setoption = MatSetOption_Preallocator; 270 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 271 272 /* special MATPREALLOCATOR functions */ 273 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator)); 274 PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATPREALLOCATOR)); 275 PetscFunctionReturn(0); 276 } 277