1 /* 2 used by MPIAIJ, BAIJ and SBAIJ to reduce code duplication 3 4 define TYPE to AIJ BAIJ or SBAIJ 5 TYPE_SBAIJ for SBAIJ matrix 6 7 */ 8 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 16 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 } 51 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; 56 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 } 64 65 static PetscErrorCode MatFinishScatterAndSetValues_MPI_Hash(Mat A) 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; 72 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; 78 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 PetscFunctionReturn(PETSC_SUCCESS); 94 } 95 96 static PetscErrorCode MatAssemblyEnd_MPI_Hash(Mat A, MatAssemblyType type) 97 { 98 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 99 100 PetscFunctionBegin; 101 PetscCall(MatFinishScatterAndSetValues_MPI_Hash(A)); 102 if (type != MAT_FINAL_ASSEMBLY) PetscFunctionReturn(PETSC_SUCCESS); 103 104 A->insertmode = NOT_SET_VALUES; /* this was set by the previous calls to MatSetValues() */ 105 106 A->ops[0] = a->cops; 107 A->hash_active = PETSC_FALSE; 108 109 PetscCall(MatAssemblyBegin(a->A, MAT_FINAL_ASSEMBLY)); 110 PetscCall(MatAssemblyEnd(a->A, MAT_FINAL_ASSEMBLY)); 111 PetscCall(MatAssemblyBegin(a->B, MAT_FINAL_ASSEMBLY)); 112 PetscCall(MatAssemblyEnd(a->B, MAT_FINAL_ASSEMBLY)); 113 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 114 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 115 PetscFunctionReturn(PETSC_SUCCESS); 116 } 117 118 static PetscErrorCode MatCopyHashToXAIJ_MPI_Hash(Mat A, Mat B) 119 { 120 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data, *b = (PetscConcat(Mat_MPI, TYPE) *)B->data; 121 122 PetscFunctionBegin; 123 /* Let's figure there's no harm done in doing the scatters for A now even if A != B */ 124 PetscCall(MatAssemblyBegin_MPI_Hash(A, /*unused*/ MAT_FINAL_ASSEMBLY)); 125 PetscCall(MatFinishScatterAndSetValues_MPI_Hash(A)); 126 127 PetscCall(MatCopyHashToXAIJ(a->A, b->A)); 128 PetscCall(MatCopyHashToXAIJ(a->B, b->B)); 129 PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY)); 130 PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY)); 131 PetscFunctionReturn(PETSC_SUCCESS); 132 } 133 134 static PetscErrorCode MatDestroy_MPI_Hash(Mat A) 135 { 136 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 137 138 PetscFunctionBegin; 139 PetscCall(MatStashDestroy_Private(&A->stash)); 140 PetscCall(MatDestroy(&a->A)); 141 PetscCall(MatDestroy(&a->B)); 142 PetscCall((*a->cops.destroy)(A)); 143 PetscFunctionReturn(PETSC_SUCCESS); 144 } 145 146 static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A) 147 { 148 PetscFunctionBegin; 149 PetscFunctionReturn(PETSC_SUCCESS); 150 } 151 152 static PetscErrorCode MatSetRandom_MPI_Hash(Mat A, PETSC_UNUSED PetscRandom r) 153 { 154 SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Must set preallocation first"); 155 } 156 157 static PetscErrorCode MatSetUp_MPI_Hash(Mat A) 158 { 159 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 160 PetscMPIInt size; 161 #if !defined(TYPE_AIJ) 162 PetscInt bs; 163 #endif 164 165 PetscFunctionBegin; 166 PetscCall(PetscInfo(A, "Using hash-based MatSetValues() for MATMPISBAIJ because no preallocation provided\n")); 167 PetscCall(PetscLayoutSetUp(A->rmap)); 168 PetscCall(PetscLayoutSetUp(A->cmap)); 169 if (A->rmap->bs < 1) A->rmap->bs = 1; 170 if (A->cmap->bs < 1) A->cmap->bs = 1; 171 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size)); 172 173 #if !defined(TYPE_AIJ) 174 PetscCall(MatGetBlockSize(A, &bs)); 175 /* these values are set in MatMPISBAIJSetPreallocation() */ 176 a->bs2 = bs * bs; 177 a->mbs = A->rmap->n / bs; 178 a->nbs = A->cmap->n / bs; 179 a->Mbs = A->rmap->N / bs; 180 a->Nbs = A->cmap->N / bs; 181 182 for (PetscInt i = 0; i <= a->size; i++) a->rangebs[i] = A->rmap->range[i] / bs; 183 a->rstartbs = A->rmap->rstart / bs; 184 a->rendbs = A->rmap->rend / bs; 185 a->cstartbs = A->cmap->rstart / bs; 186 a->cendbs = A->cmap->rend / bs; 187 PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), A->rmap->bs, &A->bstash)); 188 #endif 189 190 PetscCall(MatCreate(PETSC_COMM_SELF, &a->A)); 191 PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->n, A->rmap->n, A->cmap->n)); 192 PetscCall(MatSetBlockSizesFromMats(a->A, A, A)); 193 #if defined(SUB_TYPE_CUSPARSE) 194 PetscCall(MatSetType(a->A, MATSEQAIJCUSPARSE)); 195 #else 196 PetscCall(MatSetType(a->A, PetscConcat(MATSEQ, TYPE))); 197 #endif 198 PetscCall(MatSetUp(a->A)); 199 200 PetscCall(MatCreate(PETSC_COMM_SELF, &a->B)); 201 PetscCall(MatSetSizes(a->B, A->rmap->n, size > 1 ? A->cmap->N : 0, A->rmap->n, size > 1 ? A->cmap->N : 0)); 202 PetscCall(MatSetBlockSizesFromMats(a->B, A, A)); 203 #if defined(TYPE_SBAIJ) 204 PetscCall(MatSetType(a->B, MATSEQBAIJ)); 205 #else 206 #if defined(SUB_TYPE_CUSPARSE) 207 PetscCall(MatSetType(a->B, MATSEQAIJCUSPARSE)); 208 #else 209 PetscCall(MatSetType(a->B, PetscConcat(MATSEQ, TYPE))); 210 #endif 211 #endif 212 PetscCall(MatSetUp(a->B)); 213 214 /* keep a record of the operations so they can be reset when the hash handling is complete */ 215 a->cops = A->ops[0]; 216 A->ops->assemblybegin = MatAssemblyBegin_MPI_Hash; 217 A->ops->assemblyend = MatAssemblyEnd_MPI_Hash; 218 A->ops->setvalues = MatSetValues_MPI_Hash; 219 A->ops->destroy = MatDestroy_MPI_Hash; 220 A->ops->zeroentries = MatZeroEntries_MPI_Hash; 221 A->ops->setrandom = MatSetRandom_MPI_Hash; 222 A->ops->copyhashtoxaij = MatCopyHashToXAIJ_MPI_Hash; 223 A->ops->setvaluesblocked = NULL; 224 225 A->preallocated = PETSC_TRUE; 226 A->hash_active = PETSC_TRUE; 227 PetscFunctionReturn(PETSC_SUCCESS); 228 } 229