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