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 ierr = MatStashCreate_Private(PetscObjectComm((PetscObject) A), bs, &A->stash);CHKERRQ(ierr); 38 /* arrays are for blocked rows/cols */ 39 mbs = m/bs; 40 ierr = PetscCalloc4(mbs, &p->dnz, mbs, &p->onz, mbs, &p->dnzu, mbs, &p->onzu);CHKERRQ(ierr); 41 PetscFunctionReturn(0); 42 } 43 44 PetscErrorCode MatSetValues_Preallocator(Mat A, PetscInt m, const PetscInt *rows, PetscInt n, const PetscInt *cols, const PetscScalar *values, InsertMode addv) 45 { 46 Mat_Preallocator *p = (Mat_Preallocator *) A->data; 47 PetscInt rStart, rEnd, r, cStart, cEnd, c, bs; 48 PetscErrorCode ierr; 49 50 PetscFunctionBegin; 51 ierr = MatGetBlockSize(A, &bs);CHKERRQ(ierr); 52 ierr = MatGetOwnershipRange(A, &rStart, &rEnd);CHKERRQ(ierr); 53 ierr = MatGetOwnershipRangeColumn(A, &cStart, &cEnd);CHKERRQ(ierr); 54 for (r = 0; r < m; ++r) { 55 PetscHashIJKey key; 56 PetscBool missing; 57 58 key.i = rows[r]; 59 if (key.i < 0) continue; 60 if ((key.i < rStart) || (key.i >= rEnd)) { 61 ierr = MatStashValuesRow_Private(&A->stash, key.i, n, cols, values, PETSC_FALSE);CHKERRQ(ierr); 62 } else { /* Hash table is for blocked rows/cols */ 63 key.i = rows[r]/bs; 64 for (c = 0; c < n; ++c) { 65 key.j = cols[c]/bs; 66 if (key.j < 0) continue; 67 ierr = PetscHSetIJQueryAdd(p->ht, key, &missing);CHKERRQ(ierr); 68 if (missing) { 69 if ((key.j >= cStart/bs) && (key.j < cEnd/bs)) { 70 ++p->dnz[key.i-rStart/bs]; 71 if (key.j >= key.i) ++p->dnzu[key.i-rStart/bs]; 72 } else { 73 ++p->onz[key.i-rStart/bs]; 74 if (key.j >= key.i) ++p->onzu[key.i-rStart/bs]; 75 } 76 } 77 } 78 } 79 } 80 PetscFunctionReturn(0); 81 } 82 83 PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type) 84 { 85 PetscInt nstash, reallocs; 86 PetscErrorCode ierr; 87 88 PetscFunctionBegin; 89 ierr = MatStashScatterBegin_Private(A, &A->stash, A->rmap->range);CHKERRQ(ierr); 90 ierr = MatStashGetInfo_Private(&A->stash, &nstash, &reallocs);CHKERRQ(ierr); 91 ierr = PetscInfo2(A, "Stash has %D entries, uses %D mallocs.\n", nstash, reallocs);CHKERRQ(ierr); 92 PetscFunctionReturn(0); 93 } 94 95 PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type) 96 { 97 PetscScalar *val; 98 PetscInt *row, *col; 99 PetscInt i, j, rstart, ncols, flg; 100 PetscMPIInt n; 101 PetscErrorCode ierr; 102 103 PetscFunctionBegin; 104 while (1) { 105 ierr = MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg);CHKERRQ(ierr); 106 if (!flg) break; 107 108 for (i = 0; i < n; ) { 109 /* Now identify the consecutive vals belonging to the same row */ 110 for (j = i, rstart = row[j]; j < n; j++) { 111 if (row[j] != rstart) break; 112 } 113 if (j < n) ncols = j-i; 114 else ncols = n-i; 115 /* Now assemble all these values with a single function call */ 116 ierr = MatSetValues_Preallocator(A, 1, row+i, ncols, col+i, val+i, INSERT_VALUES);CHKERRQ(ierr); 117 i = j; 118 } 119 } 120 ierr = MatStashScatterEnd_Private(&A->stash);CHKERRQ(ierr); 121 PetscFunctionReturn(0); 122 } 123 124 PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer) 125 { 126 PetscFunctionBegin; 127 PetscFunctionReturn(0); 128 } 129 130 PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg) 131 { 132 PetscFunctionBegin; 133 PetscFunctionReturn(0); 134 } 135 136 PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A) 137 { 138 Mat_Preallocator *p = (Mat_Preallocator *) mat->data; 139 PetscInt bs; 140 PetscErrorCode ierr; 141 142 PetscFunctionBegin; 143 ierr = MatGetBlockSize(mat, &bs);CHKERRQ(ierr); 144 ierr = MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);CHKERRQ(ierr); 145 ierr = MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);CHKERRQ(ierr); 146 PetscFunctionReturn(0); 147 } 148 149 /*@ 150 MatPreallocatorPreallocate - Preallocates the input matrix, optionally filling it with zeros 151 152 Input Parameter: 153 + mat - the preallocator 154 - fill - fill the matrix with zeros 155 156 Output Parameter: 157 . A - the matrix 158 159 Level: advanced 160 161 .seealso: MATPREALLOCATOR 162 @*/ 163 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A) 164 { 165 PetscErrorCode ierr; 166 167 PetscFunctionBegin; 168 PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 169 PetscValidHeaderSpecific(A, MAT_CLASSID, 3); 170 ierr = PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));CHKERRQ(ierr); 171 PetscFunctionReturn(0); 172 } 173 174 /*MC 175 MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation. 176 177 Operations Provided: 178 . MatSetValues() 179 180 Options Database Keys: 181 . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions() 182 183 Level: advanced 184 185 .seealso: Mat 186 187 M*/ 188 189 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A) 190 { 191 Mat_Preallocator *p; 192 PetscErrorCode ierr; 193 194 PetscFunctionBegin; 195 ierr = PetscNewLog(A, &p);CHKERRQ(ierr); 196 A->data = (void *) p; 197 198 p->ht = NULL; 199 p->dnz = NULL; 200 p->onz = NULL; 201 p->dnzu = NULL; 202 p->onzu = NULL; 203 204 /* matrix ops */ 205 ierr = PetscMemzero(A->ops, sizeof(struct _MatOps));CHKERRQ(ierr); 206 207 A->ops->destroy = MatDestroy_Preallocator; 208 A->ops->setup = MatSetUp_Preallocator; 209 A->ops->setvalues = MatSetValues_Preallocator; 210 A->ops->assemblybegin = MatAssemblyBegin_Preallocator; 211 A->ops->assemblyend = MatAssemblyEnd_Preallocator; 212 A->ops->view = MatView_Preallocator; 213 A->ops->setoption = MatSetOption_Preallocator; 214 A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */ 215 216 /* special MATPREALLOCATOR functions */ 217 ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);CHKERRQ(ierr); 218 ierr = PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);CHKERRQ(ierr); 219 PetscFunctionReturn(0); 220 } 221