Date: 2025-03-28T21:10:58+00:00

Actual source code: mpihashmat.h

      1: /*
      2:    used by MPIAIJ, BAIJ and SBAIJ to reduce code duplication

      4:      define TYPE to AIJ BAIJ or SBAIJ
      5:             TYPE_SBAIJ for SBAIJ matrix

      7: */

      9: static PetscErrorCode MatSetValues_MPI_Hash(Mat A, PetscInt m, const PetscInt *rows, PetscInt n, const PetscInt *cols, const PetscScalar *values, InsertMode addv)
     10: {
     11:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
     12:   PetscInt rStart, rEnd, cStart, cEnd;
     13: #if defined(TYPE_SBAIJ)
     14:   PetscInt bs;
     15: #endif

     17:   PetscFunctionBegin;
     18:   PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd));
     19:   PetscCall(MatGetOwnershipRangeColumn(A, &cStart, &cEnd));
     20: #if defined(TYPE_SBAIJ)
     21:   PetscCall(MatGetBlockSize(A, &bs));
     22: #endif
     23:   for (PetscInt r = 0; r < m; ++r) {
     24:     PetscScalar value;
     25:     if (rows[r] < 0) continue;
     26:     if (rows[r] < rStart || rows[r] >= rEnd) {
     27:       PetscCheck(!A->nooffprocentries, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Setting off process row %" PetscInt_FMT " even though MatSetOption(,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE) was set", rows[r]);
     28:       if (!a->donotstash) {
     29:         A->assembled = PETSC_FALSE;
     30:         if (a->roworiented) {
     31:           PetscCall(MatStashValuesRow_Private(&A->stash, rows[r], n, cols, values + r * n, PETSC_FALSE));
     32:         } else {
     33:           PetscCall(MatStashValuesCol_Private(&A->stash, rows[r], n, cols, values + r, m, PETSC_FALSE));
     34:         }
     35:       }
     36:     } else {
     37:       for (PetscInt c = 0; c < n; ++c) {
     38: #if defined(TYPE_SBAIJ)
     39:         if (cols[c] / bs < rows[r] / bs) continue;
     40: #else
     41:         if (cols[c] < 0) continue;
     42: #endif
     43:         value = values ? (a->roworiented ? values[r * n + c] : values[r + m * c]) : 0;
     44:         if (cols[c] >= cStart && cols[c] < cEnd) PetscCall(MatSetValue(a->A, rows[r] - rStart, cols[c] - cStart, value, addv));
     45:         else PetscCall(MatSetValue(a->B, rows[r] - rStart, cols[c], value, addv));
     46:       }
     47:     }
     48:   }
     49:   PetscFunctionReturn(PETSC_SUCCESS);
     50: }

     52: static PetscErrorCode MatAssemblyBegin_MPI_Hash(Mat A, PETSC_UNUSED MatAssemblyType type)
     53: {
     54:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
     55:   PetscInt nstash, reallocs;

     57:   PetscFunctionBegin;
     58:   if (a->donotstash || A->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);
     59:   PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range));
     60:   PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs));
     61:   PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
     62:   PetscFunctionReturn(PETSC_SUCCESS);
     63: }

     65: static PetscErrorCode MatAssemblyEnd_MPI_Hash(Mat A, MatAssemblyType type)
     66: {
     67:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
     68:   PetscMPIInt  n;
     69:   PetscScalar *val;
     70:   PetscInt    *row, *col;
     71:   PetscInt     j, ncols, flg, rstart;

     73:   PetscFunctionBegin;
     74:   if (!a->donotstash && !A->nooffprocentries) {
     75:     while (1) {
     76:       PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg));
     77:       if (!flg) break;

     79:       for (PetscInt i = 0; i < n;) {
     80:         /* Now identify the consecutive vals belonging to the same row */
     81:         for (j = i, rstart = row[j]; j < n; j++) {
     82:           if (row[j] != rstart) break;
     83:         }
     84:         if (j < n) ncols = j - i;
     85:         else ncols = n - i;
     86:         /* Now assemble all these values with a single function call */
     87:         PetscCall(MatSetValues_MPI_Hash(A, 1, row + i, ncols, col + i, val + i, A->insertmode));
     88:         i = j;
     89:       }
     90:     }
     91:     PetscCall(MatStashScatterEnd_Private(&A->stash));
     92:   }
     93:   if (type != MAT_FINAL_ASSEMBLY) PetscFunctionReturn(PETSC_SUCCESS);

     95:   A->insertmode = NOT_SET_VALUES; /* this was set by the previous calls to MatSetValues() */

     97:   A->ops[0]      = a->cops;
     98:   A->hash_active = PETSC_FALSE;

    100:   PetscCall(MatAssemblyBegin(a->A, MAT_FINAL_ASSEMBLY));
    101:   PetscCall(MatAssemblyEnd(a->A, MAT_FINAL_ASSEMBLY));
    102:   PetscCall(MatAssemblyBegin(a->B, MAT_FINAL_ASSEMBLY));
    103:   PetscCall(MatAssemblyEnd(a->B, MAT_FINAL_ASSEMBLY));
    104:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    105:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    106:   PetscFunctionReturn(PETSC_SUCCESS);
    107: }

    109: static PetscErrorCode MatDestroy_MPI_Hash(Mat A)
    110: {
    111:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;

    113:   PetscFunctionBegin;
    114:   PetscCall(MatStashDestroy_Private(&A->stash));
    115:   PetscCall(MatDestroy(&a->A));
    116:   PetscCall(MatDestroy(&a->B));
    117:   PetscCall((*a->cops.destroy)(A));
    118:   PetscFunctionReturn(PETSC_SUCCESS);
    119: }

    121: static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A)
    122: {
    123:   PetscFunctionBegin;
    124:   PetscFunctionReturn(PETSC_SUCCESS);
    125: }

    127: static PetscErrorCode MatSetRandom_MPI_Hash(Mat A, PETSC_UNUSED PetscRandom r)
    128: {
    129:   SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Must set preallocation first");
    130: }

    132: static PetscErrorCode MatSetUp_MPI_Hash(Mat A)
    133: {
    134:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
    135:   PetscMPIInt size;
    136: #if !defined(TYPE_AIJ)
    137:   PetscInt bs;
    138: #endif

    140:   PetscFunctionBegin;
    141:   PetscCall(PetscInfo(A, "Using hash-based MatSetValues() for MATMPISBAIJ because no preallocation provided\n"));
    142:   PetscCall(PetscLayoutSetUp(A->rmap));
    143:   PetscCall(PetscLayoutSetUp(A->cmap));
    144:   if (A->rmap->bs < 1) A->rmap->bs = 1;
    145:   if (A->cmap->bs < 1) A->cmap->bs = 1;
    146:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));

    148: #if !defined(TYPE_AIJ)
    149:   PetscCall(MatGetBlockSize(A, &bs));
    150:   /* these values are set in MatMPISBAIJSetPreallocation() */
    151:   a->bs2 = bs * bs;
    152:   a->mbs = A->rmap->n / bs;
    153:   a->nbs = A->cmap->n / bs;
    154:   a->Mbs = A->rmap->N / bs;
    155:   a->Nbs = A->cmap->N / bs;

    157:   for (PetscInt i = 0; i <= a->size; i++) a->rangebs[i] = A->rmap->range[i] / bs;
    158:   a->rstartbs = A->rmap->rstart / bs;
    159:   a->rendbs   = A->rmap->rend / bs;
    160:   a->cstartbs = A->cmap->rstart / bs;
    161:   a->cendbs   = A->cmap->rend / bs;
    162:   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), A->rmap->bs, &A->bstash));
    163: #endif

    165:   PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
    166:   PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->n, A->rmap->n, A->cmap->n));
    167:   PetscCall(MatSetBlockSizesFromMats(a->A, A, A));
    168: #if defined(SUB_TYPE_CUSPARSE)
    169:   PetscCall(MatSetType(a->A, MATSEQAIJCUSPARSE));
    170: #else
    171:   PetscCall(MatSetType(a->A, PetscConcat(MATSEQ, TYPE)));
    172: #endif
    173:   PetscCall(MatSetUp(a->A));

    175:   PetscCall(MatCreate(PETSC_COMM_SELF, &a->B));
    176:   PetscCall(MatSetSizes(a->B, A->rmap->n, size > 1 ? A->cmap->N : 0, A->rmap->n, size > 1 ? A->cmap->N : 0));
    177:   PetscCall(MatSetBlockSizesFromMats(a->B, A, A));
    178: #if defined(TYPE_SBAIJ)
    179:   PetscCall(MatSetType(a->B, MATSEQBAIJ));
    180: #else
    181:   #if defined(SUB_TYPE_CUSPARSE)
    182:   PetscCall(MatSetType(a->B, MATSEQAIJCUSPARSE));
    183:   #else
    184:   PetscCall(MatSetType(a->B, PetscConcat(MATSEQ, TYPE)));
    185:   #endif
    186: #endif
    187:   PetscCall(MatSetUp(a->B));

    189:   /* keep a record of the operations so they can be reset when the hash handling is complete */
    190:   a->cops                  = A->ops[0];
    191:   A->ops->assemblybegin    = MatAssemblyBegin_MPI_Hash;
    192:   A->ops->assemblyend      = MatAssemblyEnd_MPI_Hash;
    193:   A->ops->setvalues        = MatSetValues_MPI_Hash;
    194:   A->ops->destroy          = MatDestroy_MPI_Hash;
    195:   A->ops->zeroentries      = MatZeroEntries_MPI_Hash;
    196:   A->ops->setrandom        = MatSetRandom_MPI_Hash;
    197:   A->ops->setvaluesblocked = NULL;

    199:   A->preallocated = PETSC_TRUE;
    200:   A->hash_active  = PETSC_TRUE;
    201:   PetscFunctionReturn(PETSC_SUCCESS);
    202: }