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

  4: typedef struct {
  5:   PetscHSetIJ 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:   PetscHSetIJDestroy(&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:   PetscHSetIJCreate(&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:     PetscHashIJKey key;
 53:     PetscBool      missing;

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

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

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

 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;

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

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

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

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

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;

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

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

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

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

152:   Level: advanced

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

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

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

170:    Operations Provided:
171: .  MatSetValues()

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

176:   Level: advanced

178: .seealso: Mat

180: M*/

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

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

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

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

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