1:  #include <petsc/private/matimpl.h>
  2:  #include <petsc/private/hash.h>

  4: typedef struct {
  5:   PetscHashJK ht;
  6:   PetscInt   *dnz, *onz;
  7: } Mat_Preallocator;

  9: PetscErrorCode MatDestroy_Preallocator(Mat A)
 10: {
 11:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
 12:   PetscErrorCode    ierr;

 15:   MatStashDestroy_Private(&A->stash);
 16:   PetscHashJKDestroy(&p->ht);
 17:   PetscFree2(p->dnz, p->onz);
 18:   PetscFree(A->data);
 19:   PetscObjectChangeTypeName((PetscObject) A, 0);
 20:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", NULL);
 21:   return(0);
 22: }

 24: PetscErrorCode MatSetUp_Preallocator(Mat A)
 25: {
 26:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
 27:   PetscInt          m, bs;
 28:   PetscErrorCode    ierr;

 31:   PetscLayoutSetUp(A->rmap);
 32:   PetscLayoutSetUp(A->cmap);
 33:   MatGetLocalSize(A, &m, NULL);
 34:   PetscHashJKCreate(&p->ht);
 35:   MatGetBlockSize(A, &bs);
 36:   MatStashCreate_Private(PetscObjectComm((PetscObject) A), bs, &A->stash);
 37:   PetscCalloc2(m, &p->dnz, m, &p->onz);
 38:   return(0);
 39: }

 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;

 48:   /* TODO: Handle blocksize */
 49:   MatGetOwnershipRange(A, &rStart, &rEnd);
 50:   MatGetOwnershipRangeColumn(A, &cStart, &cEnd);
 51:   for (r = 0; r < m; ++r) {
 52:     PetscHashJKKey  key;
 53:     PetscHashJKIter missing, iter;

 55:     key.j = rows[r];
 56:     if (key.j < 0) continue;
 57:     if ((key.j < rStart) || (key.j >= rEnd)) {
 58:       MatStashValuesRow_Private(&A->stash, key.j, n, cols, values, PETSC_FALSE);
 59:     } else {
 60:       for (c = 0; c < n; ++c) {
 61:         key.k = cols[c];
 62:         if (key.k < 0) continue;
 63:         PetscHashJKPut(p->ht, key, &missing, &iter);
 64:         if (missing) {
 65:           PetscHashJKSet(p->ht, iter, 1);
 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:   return(0);
 73: }

 75: PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
 76: {
 77:   PetscInt       nstash, reallocs;

 81:   PetscLayoutSetUp(A->rmap);
 82:   MatStashScatterBegin_Private(A, &A->stash, A->rmap->range);
 83:   MatStashGetInfo_Private(&A->stash, &nstash, &reallocs);
 84:   PetscInfo2(A, "Stash has %D entries, uses %D mallocs.\n", nstash, reallocs);
 85:   return(0);
 86: }

 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;

 97:   while (1) {
 98:     MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg);
 99:     if (!flg) break;

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:       MatSetValues_Preallocator(A, 1, row+i, ncols, col+i, val+i, INSERT_VALUES);
110:       i = j;
111:     }
112:   }
113:   MatStashScatterEnd_Private(&A->stash);
114:   return(0);
115: }

117: PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
118: {
120:   return(0);
121: }

123: PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
124: {
126:   return(0);
127: }

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;

137:   MatGetBlockSize(mat, &bs);
138:   MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, udnz, uonz);
139:   MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
140:   return(0);
141: }

143: /*@
144:   MatPreallocatorPreallocate - Preallocates the input matrix, optionally filling it with zeros

146:   Input Parameter:
147: + mat  - the preallocator
148: - fill - fill the matrix with zeros

150:   Output Parameter:
151: . A    - the matrix

153:   Level: advanced

155: .seealso: MATPREALLOCATOR
156: @*/
157: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
158: {

164:   PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));
165:   return(0);
166: }

168: /*MC
169:    MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation.

171:    Operations Provided:
172: .  MatSetValues()

174:    Options Database Keys:
175: . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions()

177:   Level: advanced

179: .seealso: Mat

181: M*/

183: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
184: {
185:   Mat_Preallocator *p;
186:   PetscErrorCode    ierr;

189:   PetscNewLog(A, &p);
190:   A->data = (void *) p;

192:   p->ht  = NULL;
193:   p->dnz = NULL;
194:   p->onz = NULL;

196:   /* matrix ops */
197:   PetscMemzero(A->ops, sizeof(struct _MatOps));
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;

206:   /* special MATPREALLOCATOR functions */
207:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);
208:   PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);
209:   return(0);
210: }