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