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 mat, optionally filling A with zeros 192 193 Input Parameters: 194 + mat - the preallocator 195 . fill - fill the matrix with zeros 196 - A - the matrix to be preallocated 197 198 Notes: 199 This Mat 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 Mats need to be preallocated, consider using MatDuplicate() after 215 this function. 216 217 Level: advanced 218 219 .seealso: `MATPREALLOCATOR` 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 "preallocator" during a call to MatSetFromOptions() 240 241 Level: advanced 242 243 .seealso: `Mat`, `MatPreallocatorPreallocate()` 244 245 M*/ 246 247 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) { 248 Mat_Preallocator *p; 249 250 PetscFunctionBegin; 251 PetscCall(PetscNewLog(A, &p)); 252 A->data = (void *)p; 253 254 p->ht = NULL; 255 p->dnz = NULL; 256 p->onz = NULL; 257 p->dnzu = NULL; 258 p->onzu = NULL; 259 p->used = PETSC_FALSE; 260 261 /* matrix ops */ 262 PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps))); 263 264 A->ops->destroy = MatDestroy_Preallocator; 265 A->ops->setup = MatSetUp_Preallocator; 266 A->ops->setvalues = MatSetValues_Preallocator; 267 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 268 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 269 A->ops->view = MatView_Preallocator; 270 A->ops->setoption = MatSetOption_Preallocator; 271 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 272 273 /* special MATPREALLOCATOR functions */ 274 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator)); 275 PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATPREALLOCATOR)); 276 PetscFunctionReturn(0); 277 } 278