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