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