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 ierr = MatGetBlockSize(mat, &bs);CHKERRQ(ierr); 150 ierr = MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);CHKERRQ(ierr); 151 ierr = MatSetUp(A);CHKERRQ(ierr); 152 ierr = MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);CHKERRQ(ierr); 153 ierr = MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc);CHKERRQ(ierr); 154 if (fill) { 155 PetscHashIter hi; 156 PetscHashIJKey key; 157 PetscScalar *zeros; 158 PetscInt n,maxrow=1,*rows,*cols; 159 160 ierr = PetscHSetIJGetSize(p->ht,&n);CHKERRQ(ierr); 161 ierr = PetscMalloc2(n,&rows,n,&cols);CHKERRQ(ierr); 162 PetscHashIterBegin(p->ht,hi); 163 for (PetscInt i=0; !PetscHashIterAtEnd(p->ht,hi); i++) { 164 PetscHashIterGetKey(p->ht,hi,key); 165 rows[i] = key.i; 166 cols[i] = key.j; 167 PetscHashIterNext(p->ht,hi); 168 } 169 ierr = PetscHSetIJDestroy(&p->ht);CHKERRQ(ierr); 170 171 ierr = PetscCalloc1(maxrow*bs*bs,&zeros);CHKERRQ(ierr); 172 ierr = PetscSortIntWithArray(n,rows,cols);CHKERRQ(ierr); 173 for (PetscInt start=0,end=1; start<n; start=end,end++) { 174 while (end < n && rows[end] == rows[start]) end++; 175 ierr = PetscSortInt(end-start,&cols[start]);CHKERRQ(ierr); 176 if (maxrow < end-start) { 177 maxrow = 2*(end - start); 178 ierr = PetscRealloc(maxrow*bs*bs*sizeof(zeros[0]),&zeros);CHKERRQ(ierr); 179 ierr = PetscMemzero(zeros, maxrow*bs*bs*sizeof(zeros[0]));CHKERRQ(ierr); 180 } 181 ierr = MatSetValuesBlocked(A, 1, &rows[start], end-start, &cols[start], zeros, INSERT_VALUES);CHKERRQ(ierr); 182 } 183 ierr = PetscFree(zeros);CHKERRQ(ierr); 184 ierr = PetscFree2(rows,cols);CHKERRQ(ierr); 185 186 ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 187 ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 188 } 189 PetscFunctionReturn(0); 190 } 191 192 /*@ 193 MatPreallocatorPreallocate - Preallocates the A matrix, using information from mat, optionally filling A with zeros 194 195 Input Parameter: 196 + mat - the preallocator 197 . fill - fill the matrix with zeros 198 - A - the matrix to be preallocated 199 200 Notes: 201 This Mat implementation provides a helper utility to define the correct 202 preallocation data for a given nonzero structure. Use this object like a 203 regular matrix, e.g. loop over the nonzero structure of the matrix and 204 call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations. 205 The matrix entires provided to MatSetValues() will be ignored, it only uses 206 the row / col indices provided to determine the information required to be 207 passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero 208 structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat. 209 210 After you have assembled the preallocator matrix (mat), call MatPreallocatorPreallocate() 211 to define the preallocation information on the matrix (A). Setting the parameter 212 fill = PETSC_TRUE will insert zeros into the matrix A. Internally MatPreallocatorPreallocate() 213 will call MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE); 214 215 Level: advanced 216 217 .seealso: MATPREALLOCATOR 218 @*/ 219 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A) 220 { 221 PetscErrorCode ierr; 222 223 PetscFunctionBegin; 224 PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 225 PetscValidHeaderSpecific(A, MAT_CLASSID, 3); 226 ierr = PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));CHKERRQ(ierr); 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 . MatSetValues() 235 236 Options Database Keys: 237 . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions() 238 239 Level: advanced 240 241 .seealso: Mat 242 243 M*/ 244 245 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) 246 { 247 Mat_Preallocator *p; 248 PetscErrorCode ierr; 249 250 PetscFunctionBegin; 251 ierr = PetscNewLog(A, &p);CHKERRQ(ierr); 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 260 /* matrix ops */ 261 ierr = PetscMemzero(A->ops, sizeof(struct _MatOps));CHKERRQ(ierr); 262 263 A->ops->destroy = MatDestroy_Preallocator; 264 A->ops->setup = MatSetUp_Preallocator; 265 A->ops->setvalues = MatSetValues_Preallocator; 266 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 267 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 268 A->ops->view = MatView_Preallocator; 269 A->ops->setoption = MatSetOption_Preallocator; 270 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 271 272 /* special MATPREALLOCATOR functions */ 273 ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);CHKERRQ(ierr); 274 ierr = PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);CHKERRQ(ierr); 275 PetscFunctionReturn(0); 276 } 277