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 } Mat_Preallocator; 9 10 PetscErrorCode MatDestroy_Preallocator(Mat A) 11 { 12 Mat_Preallocator *p = (Mat_Preallocator *) A->data; 13 PetscErrorCode ierr; 14 15 PetscFunctionBegin; 16 ierr = MatStashDestroy_Private(&A->stash);CHKERRQ(ierr); 17 ierr = PetscHSetIJDestroy(&p->ht);CHKERRQ(ierr); 18 ierr = PetscFree4(p->dnz, p->onz, p->dnzu, p->onzu);CHKERRQ(ierr); 19 ierr = PetscFree(A->data);CHKERRQ(ierr); 20 ierr = PetscObjectChangeTypeName((PetscObject) A, 0);CHKERRQ(ierr); 21 ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", NULL);CHKERRQ(ierr); 22 PetscFunctionReturn(0); 23 } 24 25 PetscErrorCode MatSetUp_Preallocator(Mat A) 26 { 27 Mat_Preallocator *p = (Mat_Preallocator *) A->data; 28 PetscInt m, bs, mbs; 29 PetscErrorCode ierr; 30 31 PetscFunctionBegin; 32 ierr = PetscLayoutSetUp(A->rmap);CHKERRQ(ierr); 33 ierr = PetscLayoutSetUp(A->cmap);CHKERRQ(ierr); 34 ierr = MatGetLocalSize(A, &m, NULL);CHKERRQ(ierr); 35 ierr = PetscHSetIJCreate(&p->ht);CHKERRQ(ierr); 36 ierr = MatGetBlockSize(A, &bs);CHKERRQ(ierr); 37 /* Do not bother bstash since MatPreallocator does not implement MatSetValuesBlocked */ 38 ierr = MatStashCreate_Private(PetscObjectComm((PetscObject) A), 1, &A->stash);CHKERRQ(ierr); 39 /* arrays are for blocked rows/cols */ 40 mbs = m/bs; 41 ierr = PetscCalloc4(mbs, &p->dnz, mbs, &p->onz, mbs, &p->dnzu, mbs, &p->onzu);CHKERRQ(ierr); 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 PetscErrorCode ierr; 50 51 PetscFunctionBegin; 52 ierr = MatGetBlockSize(A, &bs);CHKERRQ(ierr); 53 ierr = MatGetOwnershipRange(A, &rStart, &rEnd);CHKERRQ(ierr); 54 ierr = MatGetOwnershipRangeColumn(A, &cStart, &cEnd);CHKERRQ(ierr); 55 for (r = 0; r < m; ++r) { 56 PetscHashIJKey key; 57 PetscBool missing; 58 59 key.i = rows[r]; 60 if (key.i < 0) continue; 61 if ((key.i < rStart) || (key.i >= rEnd)) { 62 ierr = MatStashValuesRow_Private(&A->stash, key.i, n, cols, values, PETSC_FALSE);CHKERRQ(ierr); 63 } else { /* Hash table is for blocked rows/cols */ 64 key.i = rows[r]/bs; 65 for (c = 0; c < n; ++c) { 66 key.j = cols[c]/bs; 67 if (key.j < 0) continue; 68 ierr = PetscHSetIJQueryAdd(p->ht, key, &missing);CHKERRQ(ierr); 69 if (missing) { 70 if ((key.j >= cStart/bs) && (key.j < cEnd/bs)) { 71 ++p->dnz[key.i-rStart/bs]; 72 if (key.j >= key.i) ++p->dnzu[key.i-rStart/bs]; 73 } else { 74 ++p->onz[key.i-rStart/bs]; 75 if (key.j >= key.i) ++p->onzu[key.i-rStart/bs]; 76 } 77 } 78 } 79 } 80 } 81 PetscFunctionReturn(0); 82 } 83 84 PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type) 85 { 86 PetscInt nstash, reallocs; 87 PetscErrorCode ierr; 88 89 PetscFunctionBegin; 90 ierr = MatStashScatterBegin_Private(A, &A->stash, A->rmap->range);CHKERRQ(ierr); 91 ierr = MatStashGetInfo_Private(&A->stash, &nstash, &reallocs);CHKERRQ(ierr); 92 ierr = PetscInfo2(A, "Stash has %D entries, uses %D mallocs.\n", nstash, reallocs);CHKERRQ(ierr); 93 PetscFunctionReturn(0); 94 } 95 96 PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type) 97 { 98 PetscScalar *val; 99 PetscInt *row, *col; 100 PetscInt i, j, rstart, ncols, flg; 101 PetscMPIInt n; 102 PetscErrorCode ierr; 103 104 PetscFunctionBegin; 105 while (1) { 106 ierr = MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg);CHKERRQ(ierr); 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 ierr = MatSetValues_Preallocator(A, 1, row+i, ncols, col+i, val+i, INSERT_VALUES);CHKERRQ(ierr); 118 i = j; 119 } 120 } 121 ierr = MatStashScatterEnd_Private(&A->stash);CHKERRQ(ierr); 122 PetscFunctionReturn(0); 123 } 124 125 PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer) 126 { 127 PetscFunctionBegin; 128 PetscFunctionReturn(0); 129 } 130 131 PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg) 132 { 133 PetscFunctionBegin; 134 PetscFunctionReturn(0); 135 } 136 137 PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A) 138 { 139 Mat_Preallocator *p = (Mat_Preallocator *) mat->data; 140 PetscInt bs; 141 PetscErrorCode ierr; 142 143 PetscFunctionBegin; 144 ierr = MatGetBlockSize(mat, &bs);CHKERRQ(ierr); 145 ierr = MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);CHKERRQ(ierr); 146 ierr = MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);CHKERRQ(ierr); 147 if (fill) { 148 PetscHashIter hi; 149 PetscHashIJKey key; 150 PetscScalar *zeros; 151 152 ierr = PetscCalloc1(bs*bs,&zeros);CHKERRQ(ierr); 153 154 PetscHashIterBegin(p->ht,hi); 155 while (!PetscHashIterAtEnd(p->ht,hi)) { 156 PetscHashIterGetKey(p->ht,hi,key); 157 PetscHashIterNext(p->ht,hi); 158 ierr = MatSetValuesBlocked(A,1,&key.i,1,&key.j,zeros,INSERT_VALUES);CHKERRQ(ierr); 159 } 160 ierr = PetscFree(zeros);CHKERRQ(ierr); 161 162 ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 163 ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 164 } 165 PetscFunctionReturn(0); 166 } 167 168 /*@ 169 MatPreallocatorPreallocate - Preallocates the A matrix, using information from mat, optionally filling A with zeros 170 171 Input Parameter: 172 + mat - the preallocator 173 . fill - fill the matrix with zeros 174 - A - the matrix to be preallocated 175 176 Notes: 177 This Mat implementation provides a helper utility to define the correct 178 preallocation data for a given nonzero structure. Use this object like a 179 regular matrix, e.g. loop over the nonzero structure of the matrix and 180 call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations. 181 The matrix entires provided to MatSetValues() will be ignored, it only uses 182 the row / col indices provided to determine the information required to be 183 passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero 184 structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat. 185 186 After you have assembled the preallocator matrix (mat), call MatPreallocatorPreallocate() 187 to define the preallocation information on the matrix (A). Setting the parameter 188 fill = PETSC_TRUE will insert zeros into the matrix A. Internally MatPreallocatorPreallocate() 189 will call MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE); 190 191 Level: advanced 192 193 .seealso: MATPREALLOCATOR 194 @*/ 195 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A) 196 { 197 PetscErrorCode ierr; 198 199 PetscFunctionBegin; 200 PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 201 PetscValidHeaderSpecific(A, MAT_CLASSID, 3); 202 ierr = PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));CHKERRQ(ierr); 203 PetscFunctionReturn(0); 204 } 205 206 /*MC 207 MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation. 208 209 Operations Provided: 210 . MatSetValues() 211 212 Options Database Keys: 213 . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions() 214 215 Level: advanced 216 217 .seealso: Mat 218 219 M*/ 220 221 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) 222 { 223 Mat_Preallocator *p; 224 PetscErrorCode ierr; 225 226 PetscFunctionBegin; 227 ierr = PetscNewLog(A, &p);CHKERRQ(ierr); 228 A->data = (void *) p; 229 230 p->ht = NULL; 231 p->dnz = NULL; 232 p->onz = NULL; 233 p->dnzu = NULL; 234 p->onzu = NULL; 235 236 /* matrix ops */ 237 ierr = PetscMemzero(A->ops, sizeof(struct _MatOps));CHKERRQ(ierr); 238 239 A->ops->destroy = MatDestroy_Preallocator; 240 A->ops->setup = MatSetUp_Preallocator; 241 A->ops->setvalues = MatSetValues_Preallocator; 242 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 243 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 244 A->ops->view = MatView_Preallocator; 245 A->ops->setoption = MatSetOption_Preallocator; 246 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 247 248 /* special MATPREALLOCATOR functions */ 249 ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);CHKERRQ(ierr); 250 ierr = PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);CHKERRQ(ierr); 251 PetscFunctionReturn(0); 252 } 253