Actual source code: matpreallocator.c

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

      4: typedef struct {
      5:   PetscHSetIJ ht;
      6:   PetscInt   *dnz, *onz;
      7:   PetscInt   *dnzu, *onzu;
      8:   PetscBool   nooffproc;
      9:   PetscBool   used;
     10: } Mat_Preallocator;

     12: static PetscErrorCode MatDestroy_Preallocator(Mat A)
     13: {
     14:   Mat_Preallocator *p = (Mat_Preallocator *)A->data;

     16:   PetscFunctionBegin;
     17:   PetscCall(MatStashDestroy_Private(&A->stash));
     18:   PetscCall(PetscHSetIJDestroy(&p->ht));
     19:   PetscCall(PetscFree4(p->dnz, p->onz, p->dnzu, p->onzu));
     20:   PetscCall(PetscFree(A->data));
     21:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, NULL));
     22:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", NULL));
     23:   PetscFunctionReturn(PETSC_SUCCESS);
     24: }

     26: static PetscErrorCode MatSetUp_Preallocator(Mat A)
     27: {
     28:   Mat_Preallocator *p = (Mat_Preallocator *)A->data;
     29:   PetscInt          m, bs, mbs;

     31:   PetscFunctionBegin;
     32:   PetscCall(PetscLayoutSetUp(A->rmap));
     33:   PetscCall(PetscLayoutSetUp(A->cmap));
     34:   PetscCall(MatGetLocalSize(A, &m, NULL));
     35:   PetscCall(PetscHSetIJCreate(&p->ht));
     36:   PetscCall(MatGetBlockSize(A, &bs));
     37:   /* Do not bother bstash since MatPreallocator does not implement MatSetValuesBlocked */
     38:   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), 1, &A->stash));
     39:   /* arrays are for blocked rows/cols */
     40:   mbs = m / bs;
     41:   PetscCall(PetscCalloc4(mbs, &p->dnz, mbs, &p->onz, mbs, &p->dnzu, mbs, &p->onzu));
     42:   PetscFunctionReturn(PETSC_SUCCESS);
     43: }

     45: static PetscErrorCode MatSetValues_Preallocator(Mat A, PetscInt m, const PetscInt *rows, PetscInt n, const PetscInt *cols, const PetscScalar *values, InsertMode addv)
     46: {
     47:   Mat_Preallocator *p = (Mat_Preallocator *)A->data;
     48:   PetscInt          rStart, rEnd, r, cStart, cEnd, c, bs;

     50:   PetscFunctionBegin;
     51:   PetscCall(MatGetBlockSize(A, &bs));
     52:   PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd));
     53:   PetscCall(MatGetOwnershipRangeColumn(A, &cStart, &cEnd));
     54:   for (r = 0; r < m; ++r) {
     55:     PetscHashIJKey key;
     56:     PetscBool      missing;

     58:     key.i = rows[r];
     59:     if (key.i < 0) continue;
     60:     if ((key.i < rStart) || (key.i >= rEnd)) {
     61:       PetscCall(MatStashValuesRow_Private(&A->stash, key.i, n, cols, values, PETSC_FALSE));
     62:     } else { /* Hash table is for blocked rows/cols */
     63:       key.i = rows[r] / bs;
     64:       for (c = 0; c < n; ++c) {
     65:         key.j = cols[c] / bs;
     66:         if (key.j < 0) continue;
     67:         PetscCall(PetscHSetIJQueryAdd(p->ht, key, &missing));
     68:         if (missing) {
     69:           if ((key.j >= cStart / bs) && (key.j < cEnd / bs)) {
     70:             ++p->dnz[key.i - rStart / bs];
     71:             if (key.j >= key.i) ++p->dnzu[key.i - rStart / bs];
     72:           } else {
     73:             ++p->onz[key.i - rStart / bs];
     74:             if (key.j >= key.i) ++p->onzu[key.i - rStart / bs];
     75:           }
     76:         }
     77:       }
     78:     }
     79:   }
     80:   PetscFunctionReturn(PETSC_SUCCESS);
     81: }

     83: static PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
     84: {
     85:   PetscInt nstash, reallocs;

     87:   PetscFunctionBegin;
     88:   PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range));
     89:   PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs));
     90:   PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
     91:   PetscFunctionReturn(PETSC_SUCCESS);
     92: }

     94: static PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type)
     95: {
     96:   PetscScalar      *val;
     97:   PetscInt         *row, *col;
     98:   PetscInt          i, j, rstart, ncols, flg;
     99:   PetscMPIInt       n;
    100:   Mat_Preallocator *p = (Mat_Preallocator *)A->data;

    102:   PetscFunctionBegin;
    103:   p->nooffproc = PETSC_TRUE;
    104:   while (1) {
    105:     PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg));
    106:     if (flg) p->nooffproc = PETSC_FALSE;
    107:     if (!flg) break;

    109:     for (i = 0; i < n;) {
    110:       /* Now identify the consecutive vals belonging to the same row */
    111:       for (j = i, rstart = row[j]; j < n; j++) {
    112:         if (row[j] != rstart) break;
    113:       }
    114:       if (j < n) ncols = j - i;
    115:       else ncols = n - i;
    116:       /* Now assemble all these values with a single function call */
    117:       PetscCall(MatSetValues_Preallocator(A, 1, row + i, ncols, col + i, val + i, INSERT_VALUES));
    118:       i = j;
    119:     }
    120:   }
    121:   PetscCall(MatStashScatterEnd_Private(&A->stash));
    122:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &p->nooffproc, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
    123:   PetscFunctionReturn(PETSC_SUCCESS);
    124: }

    126: static PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
    127: {
    128:   PetscFunctionBegin;
    129:   PetscFunctionReturn(PETSC_SUCCESS);
    130: }

    132: static PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
    133: {
    134:   PetscFunctionBegin;
    135:   PetscFunctionReturn(PETSC_SUCCESS);
    136: }

    138: static PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A)
    139: {
    140:   Mat_Preallocator *p = (Mat_Preallocator *)mat->data;
    141:   PetscInt          bs;

    143:   PetscFunctionBegin;
    144:   PetscCheck(!p->used, PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "MatPreallocatorPreallocate() can only be used once for a give MatPreallocator object. Consider using MatDuplicate() after preallocation.");
    145:   p->used = PETSC_TRUE;
    146:   if (!fill) PetscCall(PetscHSetIJDestroy(&p->ht));
    147:   PetscCall(MatGetBlockSize(mat, &bs));
    148:   PetscCall(MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu));
    149:   PetscCall(MatSetUp(A));
    150:   PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE));
    151:   if (fill) {
    152:     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc));
    153:     PetscHashIter  hi;
    154:     PetscHashIJKey key;
    155:     PetscScalar   *zeros;
    156:     PetscInt       n, maxrow = 1, *cols, rStart, rEnd, *rowstarts;

    158:     PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd));
    159:     // Ownership range is in terms of scalar entries, but we deal with blocks
    160:     rStart /= bs;
    161:     rEnd /= bs;
    162:     PetscCall(PetscHSetIJGetSize(p->ht, &n));
    163:     PetscCall(PetscMalloc2(n, &cols, rEnd - rStart + 1, &rowstarts));
    164:     rowstarts[0] = 0;
    165:     for (PetscInt i = 0; i < rEnd - rStart; i++) {
    166:       rowstarts[i + 1] = rowstarts[i] + p->dnz[i] + p->onz[i];
    167:       maxrow           = PetscMax(maxrow, p->dnz[i] + p->onz[i]);
    168:     }
    169:     PetscCheck(rowstarts[rEnd - rStart] == n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Hash claims %" PetscInt_FMT " entries, but dnz+onz counts %" PetscInt_FMT, n, rowstarts[rEnd - rStart]);

    171:     PetscHashIterBegin(p->ht, hi);
    172:     while (!PetscHashIterAtEnd(p->ht, hi)) {
    173:       PetscHashIterGetKey(p->ht, hi, key);
    174:       PetscInt lrow         = key.i - rStart;
    175:       cols[rowstarts[lrow]] = key.j;
    176:       rowstarts[lrow]++;
    177:       PetscHashIterNext(p->ht, hi);
    178:     }
    179:     PetscCall(PetscHSetIJDestroy(&p->ht));

    181:     PetscCall(PetscCalloc1(maxrow * bs * bs, &zeros));
    182:     for (PetscInt i = 0; i < rEnd - rStart; i++) {
    183:       PetscInt grow = rStart + i;
    184:       PetscInt end = rowstarts[i], start = end - p->dnz[i] - p->onz[i];
    185:       PetscCall(PetscSortInt(end - start, PetscSafePointerPlusOffset(cols, start)));
    186:       PetscCall(MatSetValuesBlocked(A, 1, &grow, end - start, PetscSafePointerPlusOffset(cols, start), zeros, INSERT_VALUES));
    187:     }
    188:     PetscCall(PetscFree(zeros));
    189:     PetscCall(PetscFree2(cols, rowstarts));

    191:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    192:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    193:     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE));
    194:   }
    195:   PetscFunctionReturn(PETSC_SUCCESS);
    196: }

    198: /*@
    199:   MatPreallocatorPreallocate - Preallocates the A matrix, using information from a `MATPREALLOCATOR` mat, optionally filling A with zeros

    201:   Input Parameters:
    202: + mat  - the `MATPREALLOCATOR` preallocator matrix
    203: . fill - fill the matrix with zeros
    204: - A    - the matrix to be preallocated

    206:   Notes:
    207:   This `MatType` implementation provides a helper utility to define the correct
    208:   preallocation data for a given nonzero structure. Use this object like a
    209:   regular matrix, e.g. loop over the nonzero structure of the matrix and
    210:   call `MatSetValues()` or `MatSetValuesBlocked()` to indicate the nonzero locations.
    211:   The matrix entries provided to `MatSetValues()` will be ignored, it only uses
    212:   the row / col indices provided to determine the information required to be
    213:   passed to `MatXAIJSetPreallocation()`. Once you have looped over the nonzero
    214:   structure, you must call `MatAssemblyBegin()`, `MatAssemblyEnd()` on mat.

    216:   After you have assembled the preallocator matrix (mat), call `MatPreallocatorPreallocate()`
    217:   to define the preallocation information on the matrix (A). Setting the parameter
    218:   fill = PETSC_TRUE will insert zeros into the matrix A. Internally `MatPreallocatorPreallocate()`
    219:   will call `MatSetOption`(A, `MAT_NEW_NONZERO_ALLOCATION_ERR`, `PETSC_TRUE`);

    221:   This function may only be called once for a given `MATPREALLOCATOR` object. If
    222:   multiple `Mat`s need to be preallocated, consider using `MatDuplicate()` after
    223:   this function.

    225:   Level: advanced

    227: .seealso: `MATPREALLOCATOR`, `MatXAIJSetPreallocation()`
    228: @*/
    229: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
    230: {
    231:   PetscFunctionBegin;
    235:   PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat, PetscBool, Mat), (mat, fill, A));
    236:   PetscFunctionReturn(PETSC_SUCCESS);
    237: }

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

    242:    Operations Provided:
    243: .vb
    244:   MatSetValues()
    245: .ve

    247:    Options Database Keys:
    248: . -mat_type preallocator - sets the matrix type to `MATPREALLOCATOR` during a call to `MatSetFromOptions()`

    250:   Level: advanced

    252: .seealso: `MATPREALLOCATOR`, `Mat`, `MatPreallocatorPreallocate()`
    253: M*/

    255: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
    256: {
    257:   Mat_Preallocator *p;

    259:   PetscFunctionBegin;
    260:   PetscCall(PetscNew(&p));
    261:   A->data = (void *)p;

    263:   p->ht   = NULL;
    264:   p->dnz  = NULL;
    265:   p->onz  = NULL;
    266:   p->dnzu = NULL;
    267:   p->onzu = NULL;
    268:   p->used = PETSC_FALSE;

    270:   /* matrix ops */
    271:   PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps)));

    273:   A->ops->destroy       = MatDestroy_Preallocator;
    274:   A->ops->setup         = MatSetUp_Preallocator;
    275:   A->ops->setvalues     = MatSetValues_Preallocator;
    276:   A->ops->assemblybegin = MatAssemblyBegin_Preallocator;
    277:   A->ops->assemblyend   = MatAssemblyEnd_Preallocator;
    278:   A->ops->view          = MatView_Preallocator;
    279:   A->ops->setoption     = MatSetOption_Preallocator;
    280:   A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */

    282:   /* special MATPREALLOCATOR functions */
    283:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator));
    284:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATPREALLOCATOR));
    285:   PetscFunctionReturn(PETSC_SUCCESS);
    286: }