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 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; 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 if (type != MAT_FINAL_ASSEMBLY) PetscFunctionReturn(PETSC_SUCCESS); 94 95 A->insertmode = NOT_SET_VALUES; /* this was set by the previous calls to MatSetValues() */ 96 97 A->ops[0] = a->cops; 98 A->hash_active = PETSC_FALSE; 99 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 } 108 109 static PetscErrorCode MatDestroy_MPI_Hash(Mat A) 110 { 111 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 112 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 } 120 121 static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A) 122 { 123 PetscFunctionBegin; 124 PetscFunctionReturn(PETSC_SUCCESS); 125 } 126 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 } 131 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 139 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)); 147 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; 156 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 164 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)); 174 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)); 188 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; 198 199 A->preallocated = PETSC_TRUE; 200 A->hash_active = PETSC_TRUE; 201 PetscFunctionReturn(PETSC_SUCCESS); 202 } 203