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:       if (a->roworiented) {
 28:         PetscCall(MatStashValuesRow_Private(&A->stash, rows[r], n, cols, values + r * n, PETSC_FALSE));
 29:       } else {
 30:         PetscCall(MatStashValuesCol_Private(&A->stash, rows[r], n, cols, values + r, m, PETSC_FALSE));
 31:       }
 32:     } else {
 33:       for (PetscInt c = 0; c < n; ++c) {
 34: #if defined(TYPE_SBAIJ)
 35:         if (cols[c] / bs < rows[r] / bs) continue;
 36: #else
 37:         if (cols[c] < 0) continue;
 38: #endif
 39:         value = values ? ((a->roworiented) ? values[r * n + c] : values[r + m * c]) : 0;
 40:         if (cols[c] >= cStart && cols[c] < cEnd) PetscCall(MatSetValue(a->A, rows[r] - rStart, cols[c] - cStart, value, addv));
 41:         else PetscCall(MatSetValue(a->B, rows[r] - rStart, cols[c], value, addv));
 42:       }
 43:     }
 44:   }
 45:   PetscFunctionReturn(PETSC_SUCCESS);
 46: }

 48: static PetscErrorCode MatAssemblyBegin_MPI_Hash(Mat A, PETSC_UNUSED MatAssemblyType type)
 49: {
 50:   PetscInt nstash, reallocs;

 52:   PetscFunctionBegin;
 53:   PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range));
 54:   PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs));
 55:   PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
 56:   PetscFunctionReturn(PETSC_SUCCESS);
 57: }

 59: static PetscErrorCode MatAssemblyEnd_MPI_Hash(Mat A, MatAssemblyType type)
 60: {
 61:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
 62:   PetscMPIInt  n;
 63:   PetscScalar *val;
 64:   PetscInt    *row, *col;
 65:   PetscInt     j, ncols, flg, rstart;

 67:   PetscFunctionBegin;
 68:   while (1) {
 69:     PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg));
 70:     if (!flg) break;

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

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

 89:   PetscCall(PetscMemcpy(&A->ops, &a->cops, sizeof(*(A->ops))));
 90:   A->hash_active = PETSC_FALSE;

 92:   PetscCall(MatAssemblyBegin(a->A, MAT_FINAL_ASSEMBLY));
 93:   PetscCall(MatAssemblyEnd(a->A, MAT_FINAL_ASSEMBLY));
 94:   PetscCall(MatAssemblyBegin(a->B, MAT_FINAL_ASSEMBLY));
 95:   PetscCall(MatAssemblyEnd(a->B, MAT_FINAL_ASSEMBLY));
 96:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
 97:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
 98:   PetscFunctionReturn(PETSC_SUCCESS);
 99: }

101: static PetscErrorCode MatDestroy_MPI_Hash(Mat A)
102: {
103:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;

105:   PetscFunctionBegin;
106:   PetscCall(MatStashDestroy_Private(&A->stash));
107:   PetscCall(MatDestroy(&a->A));
108:   PetscCall(MatDestroy(&a->B));
109:   PetscCall((*a->cops.destroy)(A));
110:   PetscFunctionReturn(PETSC_SUCCESS);
111: }

113: static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A)
114: {
115:   PetscFunctionBegin;
116:   PetscFunctionReturn(PETSC_SUCCESS);
117: }

119: static PetscErrorCode MatSetRandom_MPI_Hash(Mat A, PETSC_UNUSED PetscRandom r)
120: {
121:   SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Must set preallocation first");
122: }

124: static PetscErrorCode MatSetUp_MPI_Hash(Mat A)
125: {
126:   PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
127:   PetscMPIInt size;
128: #if !defined(TYPE_AIJ)
129:   PetscInt bs;
130: #endif

132:   PetscFunctionBegin;
133:   PetscCall(PetscInfo(A, "Using hash-based MatSetValues() for MATMPISBAIJ because no preallocation provided\n"));
134:   PetscCall(PetscLayoutSetUp(A->rmap));
135:   PetscCall(PetscLayoutSetUp(A->cmap));
136:   if (A->rmap->bs < 1) A->rmap->bs = 1;
137:   if (A->cmap->bs < 1) A->cmap->bs = 1;
138:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));

140: #if !defined(TYPE_AIJ)
141:   PetscCall(MatGetBlockSize(A, &bs));
142:   /* these values are set in MatMPISBAIJSetPreallocation() */
143:   a->bs2 = bs * bs;
144:   a->mbs = A->rmap->n / bs;
145:   a->nbs = A->cmap->n / bs;
146:   a->Mbs = A->rmap->N / bs;
147:   a->Nbs = A->cmap->N / bs;

149:   for (PetscInt i = 0; i <= a->size; i++) a->rangebs[i] = A->rmap->range[i] / bs;
150:   a->rstartbs = A->rmap->rstart / bs;
151:   a->rendbs   = A->rmap->rend / bs;
152:   a->cstartbs = A->cmap->rstart / bs;
153:   a->cendbs   = A->cmap->rend / bs;
154:   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), A->rmap->bs, &A->bstash));
155: #endif

157:   PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
158:   PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->n, A->rmap->n, A->cmap->n));
159:   PetscCall(MatSetBlockSizesFromMats(a->A, A, A));
160: #if defined(SUB_TYPE_CUSPARSE)
161:   PetscCall(MatSetType(a->A, MATSEQAIJCUSPARSE));
162: #else
163:   PetscCall(MatSetType(a->A, PetscConcat(MATSEQ, TYPE)));
164: #endif
165:   PetscCall(MatSetUp(a->A));

167:   PetscCall(MatCreate(PETSC_COMM_SELF, &a->B));
168:   PetscCall(MatSetSizes(a->B, A->rmap->n, size > 1 ? A->cmap->N : 0, A->rmap->n, size > 1 ? A->cmap->N : 0));
169:   PetscCall(MatSetBlockSizesFromMats(a->B, A, A));
170: #if defined(TYPE_SBAIJ)
171:   PetscCall(MatSetType(a->B, MATSEQBAIJ));
172: #else
173:   #if defined(SUB_TYPE_CUSPARSE)
174:   PetscCall(MatSetType(a->B, MATSEQAIJCUSPARSE));
175:   #else
176:   PetscCall(MatSetType(a->B, PetscConcat(MATSEQ, TYPE)));
177:   #endif
178: #endif
179:   PetscCall(MatSetUp(a->B));

181:   /* keep a record of the operations so they can be reset when the hash handling is complete */
182:   PetscCall(PetscMemcpy(&a->cops, &A->ops, sizeof(*(A->ops))));

184:   A->ops->assemblybegin    = MatAssemblyBegin_MPI_Hash;
185:   A->ops->assemblyend      = MatAssemblyEnd_MPI_Hash;
186:   A->ops->setvalues        = MatSetValues_MPI_Hash;
187:   A->ops->destroy          = MatDestroy_MPI_Hash;
188:   A->ops->zeroentries      = MatZeroEntries_MPI_Hash;
189:   A->ops->setrandom        = MatSetRandom_MPI_Hash;
190:   A->ops->setvaluesblocked = NULL;

192:   A->preallocated = PETSC_TRUE;
193:   A->hash_active  = PETSC_TRUE;
194:   PetscFunctionReturn(PETSC_SUCCESS);
195: }