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