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