Date: 2026-06-01T23:13:31+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:   const PetscInt rStart         = A->rmap->rstart;
     13:   const PetscInt rEnd           = A->rmap->rend;
     14:   const PetscInt cStart         = A->cmap->rstart;
     15:   const PetscInt cEnd           = A->cmap->rend;
     16: #if defined(TYPE_SBAIJ)
     17:   const PetscInt bs = A->rmap->bs;
     18: #endif
     19:   const PetscBool ignorezeroentries = ((Mat_SeqAIJ *)a->A->data)->ignorezeroentries;

     21:   PetscFunctionBegin;
     22:   for (PetscInt r = 0; r < m; ++r) {
     23:     PetscScalar value;
     24:     if (rows[r] < 0) continue;
     25:     if (rows[r] < rStart || rows[r] >= rEnd) {
     26:       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]);
     27:       if (!a->donotstash) {
     28:         A->assembled = PETSC_FALSE;
     29:         if (a->roworiented) {
     30:           PetscCall(MatStashValuesRow_Private(&A->stash, rows[r], n, cols, PetscSafePointerPlusOffset(values, r * n), (PetscBool)(ignorezeroentries && (addv == ADD_VALUES))));
     31:         } else {
     32:           PetscCall(MatStashValuesCol_Private(&A->stash, rows[r], n, cols, PetscSafePointerPlusOffset(values, r), m, (PetscBool)(ignorezeroentries && (addv == ADD_VALUES))));
     33:         }
     34:       }
     35:     } else {
     36:       for (PetscInt c = 0; c < n; ++c) {
     37: #if defined(TYPE_SBAIJ)
     38:         if (cols[c] / bs < rows[r] / bs) continue;
     39: #else
     40:         if (cols[c] < 0) continue;
     41: #endif
     42:         value = values ? (a->roworiented ? values[r * n + c] : values[r + m * c]) : 0;
     43:         if (ignorezeroentries && value == 0.0 && (addv == ADD_VALUES) && rows[r] != cols[c]) continue;
     44:         if (cols[c] >= cStart && cols[c] < cEnd) {
     45:           PetscCall(MatSetValue(a->A, rows[r] - rStart, cols[c] - cStart, value, addv));
     46:         } else if (!ignorezeroentries || value != 0.0) {
     47:           /* MPIAIJ never inserts in B if it is 0 */
     48:           PetscCall(MatSetValue(a->B, rows[r] - rStart, cols[c], value, addv));
     49:         }
     50:       }
     51:     }
     52:   }
     53:   PetscFunctionReturn(PETSC_SUCCESS);
     54: }

     56: static PetscErrorCode MatAssemblyBegin_MPI_Hash(Mat A, PETSC_UNUSED MatAssemblyType type)
     57: {
     58:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
     59:   PetscInt nstash, reallocs;

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

     69: static PetscErrorCode MatFinishScatterAndSetValues_MPI_Hash(Mat A)
     70: {
     71:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
     72:   PetscMPIInt  n;
     73:   PetscScalar *val;
     74:   PetscInt    *row, *col;
     75:   PetscInt     j, ncols, flg, rstart;

     77:   PetscFunctionBegin;
     78:   if (!a->donotstash && !A->nooffprocentries) {
     79:     while (1) {
     80:       PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg));
     81:       if (!flg) break;

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

    100: static PetscErrorCode MatAssemblyEnd_MPI_Hash(Mat A, MatAssemblyType type)
    101: {
    102:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
    103:   PetscBool nooffprocentries    = A->nooffprocentries;

    105:   PetscFunctionBegin;
    106:   PetscCall(MatFinishScatterAndSetValues_MPI_Hash(A));
    107:   if (type != MAT_FINAL_ASSEMBLY) PetscFunctionReturn(PETSC_SUCCESS);

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

    111:   A->ops[0]      = a->cops;
    112:   A->hash_active = PETSC_FALSE;

    114:   /* a->B must be in CSR before MatSetUpMultiply_MPIAIJ() runs inside MatAssemblyEnd_MPIAIJ().
    115:      a->A is assembled by MatAssemblyEnd_MPIAIJ() itself and does not need a separate call here. */
    116:   PetscCall(MatAssemblyBegin(a->B, MAT_FINAL_ASSEMBLY));
    117:   PetscCall(MatAssemblyEnd(a->B, MAT_FINAL_ASSEMBLY));
    118:   /* Off-process entries were already scattered in MatAssemblyBegin_MPI_Hash().
    119:      Suppress the stash scatter in MatAssemblyBegin_MPIAIJ() to avoid a second
    120:      PetscCommBuildTwoSidedFReq() call on an empty stash. */
    121:   A->nooffprocentries = PETSC_TRUE;
    122:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    123:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    124:   A->nooffprocentries = nooffprocentries;
    125:   PetscFunctionReturn(PETSC_SUCCESS);
    126: }

    128: static PetscErrorCode MatCopyHashToXAIJ_MPI_Hash(Mat A, Mat B)
    129: {
    130:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data, *b = (PetscConcat(Mat_MPI, TYPE) *)B->data;

    132:   PetscFunctionBegin;
    133:   /* Let's figure there's no harm done in doing the scatters for A now even if A != B */
    134:   PetscCall(MatAssemblyBegin_MPI_Hash(A, /*unused*/ MAT_FINAL_ASSEMBLY));
    135:   PetscCall(MatFinishScatterAndSetValues_MPI_Hash(A));

    137:   PetscCall(MatCopyHashToXAIJ(a->A, b->A));
    138:   PetscCall(MatCopyHashToXAIJ(a->B, b->B));
    139:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
    140:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
    141:   PetscFunctionReturn(PETSC_SUCCESS);
    142: }

    144: static PetscErrorCode MatDestroy_MPI_Hash(Mat A)
    145: {
    146:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;

    148:   PetscFunctionBegin;
    149:   PetscCall(MatStashDestroy_Private(&A->stash));
    150:   PetscCall(MatDestroy(&a->A));
    151:   PetscCall(MatDestroy(&a->B));
    152:   PetscCall((*a->cops.destroy)(A));
    153:   PetscFunctionReturn(PETSC_SUCCESS);
    154: }

    156: static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A)
    157: {
    158:   PetscFunctionBegin;
    159:   PetscFunctionReturn(PETSC_SUCCESS);
    160: }

    162: static PetscErrorCode MatSetRandom_MPI_Hash(Mat A, PETSC_UNUSED PetscRandom r)
    163: {
    164:   SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Must set preallocation first");
    165: }

    167: static PetscErrorCode MatSetUp_MPI_Hash(Mat A)
    168: {
    169:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
    170:   PetscMPIInt size;
    171: #if !defined(TYPE_AIJ)
    172:   PetscInt bs;
    173: #endif

    175:   PetscFunctionBegin;
    176:   PetscCall(PetscInfo(A, "Using hash-based MatSetValues() for MATMPI" PetscStringize(TYPE) " because no preallocation provided\n"));
    177:   PetscCall(PetscLayoutSetUp(A->rmap));
    178:   PetscCall(PetscLayoutSetUp(A->cmap));
    179:   if (A->rmap->bs < 1) A->rmap->bs = 1;
    180:   if (A->cmap->bs < 1) A->cmap->bs = 1;
    181:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));

    183: #if !defined(TYPE_AIJ)
    184:   PetscCall(MatGetBlockSize(A, &bs));
    185:   /* these values are set in MatMPISBAIJSetPreallocation() */
    186:   a->bs2 = bs * bs;
    187:   a->mbs = A->rmap->n / bs;
    188:   a->nbs = A->cmap->n / bs;
    189:   a->Mbs = A->rmap->N / bs;
    190:   a->Nbs = A->cmap->N / bs;

    192:   for (PetscInt i = 0; i <= a->size; i++) a->rangebs[i] = A->rmap->range[i] / bs;
    193:   a->rstartbs = A->rmap->rstart / bs;
    194:   a->rendbs   = A->rmap->rend / bs;
    195:   a->cstartbs = A->cmap->rstart / bs;
    196:   a->cendbs   = A->cmap->rend / bs;
    197:   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), A->rmap->bs, &A->bstash));
    198: #endif

    200:   PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
    201:   PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->n, A->rmap->n, A->cmap->n));
    202:   PetscCall(MatSetBlockSizesFromMats(a->A, A, A));
    203: #if defined(SUB_TYPE_CUSPARSE)
    204:   PetscCall(MatSetType(a->A, MATSEQAIJCUSPARSE));
    205: #else
    206:   PetscCall(MatSetType(a->A, PetscConcat(MATSEQ, TYPE)));
    207: #endif
    208:   PetscCall(MatSetUp(a->A));

    210:   PetscCall(MatCreate(PETSC_COMM_SELF, &a->B));
    211:   PetscCall(MatSetSizes(a->B, A->rmap->n, size > 1 ? A->cmap->N : 0, A->rmap->n, size > 1 ? A->cmap->N : 0));
    212:   PetscCall(MatSetBlockSizesFromMats(a->B, A, A));
    213: #if defined(TYPE_SBAIJ)
    214:   PetscCall(MatSetType(a->B, MATSEQBAIJ));
    215: #else
    216:   #if defined(SUB_TYPE_CUSPARSE)
    217:   PetscCall(MatSetType(a->B, MATSEQAIJCUSPARSE));
    218:   #else
    219:   PetscCall(MatSetType(a->B, PetscConcat(MATSEQ, TYPE)));
    220:   #endif
    221: #endif
    222:   PetscCall(MatSetUp(a->B));

    224:   /* keep a record of the operations so they can be reset when the hash handling is complete */
    225:   a->cops                  = A->ops[0];
    226:   A->ops->assemblybegin    = MatAssemblyBegin_MPI_Hash;
    227:   A->ops->assemblyend      = MatAssemblyEnd_MPI_Hash;
    228:   A->ops->setvalues        = MatSetValues_MPI_Hash;
    229:   A->ops->destroy          = MatDestroy_MPI_Hash;
    230:   A->ops->zeroentries      = MatZeroEntries_MPI_Hash;
    231:   A->ops->setrandom        = MatSetRandom_MPI_Hash;
    232:   A->ops->copyhashtoxaij   = MatCopyHashToXAIJ_MPI_Hash;
    233:   A->ops->setvaluesblocked = NULL;

    235:   A->preallocated = PETSC_TRUE;
    236:   A->hash_active  = PETSC_TRUE;
    237:   PetscFunctionReturn(PETSC_SUCCESS);
    238: }