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 if (fill) { 152 PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc)); 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 while (!PetscHashIterAtEnd(p->ht,hi)) { 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 PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE)); 194 } 195 PetscFunctionReturn(0); 196 } 197 198 /*@ 199 MatPreallocatorPreallocate - Preallocates the A matrix, using information from mat, optionally filling A with zeros 200 201 Input Parameters: 202 + mat - the preallocator 203 . fill - fill the matrix with zeros 204 - A - the matrix to be preallocated 205 206 Notes: 207 This Mat implementation provides a helper utility to define the correct 208 preallocation data for a given nonzero structure. Use this object like a 209 regular matrix, e.g. loop over the nonzero structure of the matrix and 210 call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations. 211 The matrix entries provided to MatSetValues() will be ignored, it only uses 212 the row / col indices provided to determine the information required to be 213 passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero 214 structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat. 215 216 After you have assembled the preallocator matrix (mat), call MatPreallocatorPreallocate() 217 to define the preallocation information on the matrix (A). Setting the parameter 218 fill = PETSC_TRUE will insert zeros into the matrix A. Internally MatPreallocatorPreallocate() 219 will call MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE); 220 221 This function may only be called once for a given MatPreallocator object. If 222 multiple Mats need to be preallocated, consider using MatDuplicate() after 223 this function. 224 225 Level: advanced 226 227 .seealso: `MATPREALLOCATOR` 228 @*/ 229 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A) 230 { 231 PetscFunctionBegin; 232 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 233 PetscValidLogicalCollectiveBool(mat,fill,2); 234 PetscValidHeaderSpecific(A,MAT_CLASSID,3); 235 PetscUseMethod(mat,"MatPreallocatorPreallocate_C",(Mat,PetscBool,Mat),(mat,fill,A)); 236 PetscFunctionReturn(0); 237 } 238 239 /*MC 240 MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation. 241 242 Operations Provided: 243 .vb 244 MatSetValues() 245 .ve 246 247 Options Database Keys: 248 . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions() 249 250 Level: advanced 251 252 .seealso: `Mat`, `MatPreallocatorPreallocate()` 253 254 M*/ 255 256 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) 257 { 258 Mat_Preallocator *p; 259 260 PetscFunctionBegin; 261 PetscCall(PetscNewLog(A, &p)); 262 A->data = (void *) p; 263 264 p->ht = NULL; 265 p->dnz = NULL; 266 p->onz = NULL; 267 p->dnzu = NULL; 268 p->onzu = NULL; 269 p->used = PETSC_FALSE; 270 271 /* matrix ops */ 272 PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps))); 273 274 A->ops->destroy = MatDestroy_Preallocator; 275 A->ops->setup = MatSetUp_Preallocator; 276 A->ops->setvalues = MatSetValues_Preallocator; 277 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 278 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 279 A->ops->view = MatView_Preallocator; 280 A->ops->setoption = MatSetOption_Preallocator; 281 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 282 283 /* special MATPREALLOCATOR functions */ 284 PetscCall(PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator)); 285 PetscCall(PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR)); 286 PetscFunctionReturn(0); 287 } 288