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));CHKERRQ(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 159 ierr = PetscCalloc1(bs*bs,&zeros);CHKERRQ(ierr); 160 161 PetscHashIterBegin(p->ht,hi); 162 while (!PetscHashIterAtEnd(p->ht,hi)) { 163 PetscHashIterGetKey(p->ht,hi,key); 164 PetscHashIterNext(p->ht,hi); 165 ierr = MatSetValuesBlocked(A,1,&key.i,1,&key.j,zeros,INSERT_VALUES);CHKERRQ(ierr); 166 } 167 ierr = PetscFree(zeros);CHKERRQ(ierr); 168 169 ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 170 ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 171 } 172 PetscFunctionReturn(0); 173 } 174 175 /*@ 176 MatPreallocatorPreallocate - Preallocates the A matrix, using information from mat, optionally filling A with zeros 177 178 Input Parameter: 179 + mat - the preallocator 180 . fill - fill the matrix with zeros 181 - A - the matrix to be preallocated 182 183 Notes: 184 This Mat implementation provides a helper utility to define the correct 185 preallocation data for a given nonzero structure. Use this object like a 186 regular matrix, e.g. loop over the nonzero structure of the matrix and 187 call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations. 188 The matrix entires provided to MatSetValues() will be ignored, it only uses 189 the row / col indices provided to determine the information required to be 190 passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero 191 structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat. 192 193 After you have assembled the preallocator matrix (mat), call MatPreallocatorPreallocate() 194 to define the preallocation information on the matrix (A). Setting the parameter 195 fill = PETSC_TRUE will insert zeros into the matrix A. Internally MatPreallocatorPreallocate() 196 will call MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE); 197 198 Level: advanced 199 200 .seealso: MATPREALLOCATOR 201 @*/ 202 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A) 203 { 204 PetscErrorCode ierr; 205 206 PetscFunctionBegin; 207 PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 208 PetscValidHeaderSpecific(A, MAT_CLASSID, 3); 209 ierr = PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));CHKERRQ(ierr); 210 PetscFunctionReturn(0); 211 } 212 213 /*MC 214 MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation. 215 216 Operations Provided: 217 . MatSetValues() 218 219 Options Database Keys: 220 . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions() 221 222 Level: advanced 223 224 .seealso: Mat 225 226 M*/ 227 228 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) 229 { 230 Mat_Preallocator *p; 231 PetscErrorCode ierr; 232 233 PetscFunctionBegin; 234 ierr = PetscNewLog(A, &p);CHKERRQ(ierr); 235 A->data = (void *) p; 236 237 p->ht = NULL; 238 p->dnz = NULL; 239 p->onz = NULL; 240 p->dnzu = NULL; 241 p->onzu = NULL; 242 243 /* matrix ops */ 244 ierr = PetscMemzero(A->ops, sizeof(struct _MatOps));CHKERRQ(ierr); 245 246 A->ops->destroy = MatDestroy_Preallocator; 247 A->ops->setup = MatSetUp_Preallocator; 248 A->ops->setvalues = MatSetValues_Preallocator; 249 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 250 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 251 A->ops->view = MatView_Preallocator; 252 A->ops->setoption = MatSetOption_Preallocator; 253 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 254 255 /* special MATPREALLOCATOR functions */ 256 ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);CHKERRQ(ierr); 257 ierr = PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);CHKERRQ(ierr); 258 PetscFunctionReturn(0); 259 } 260