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: }