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