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 PetscInt nstash, reallocs; 55 56 PetscFunctionBegin; 57 PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range)); 58 PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs)); 59 PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs)); 60 PetscFunctionReturn(PETSC_SUCCESS); 61 } 62 63 static PetscErrorCode MatAssemblyEnd_MPI_Hash(Mat A, MatAssemblyType type) 64 { 65 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 66 PetscMPIInt n; 67 PetscScalar *val; 68 PetscInt *row, *col; 69 PetscInt j, ncols, flg, rstart; 70 71 PetscFunctionBegin; 72 while (1) { 73 PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg)); 74 if (!flg) break; 75 76 for (PetscInt i = 0; i < n;) { 77 /* Now identify the consecutive vals belonging to the same row */ 78 for (j = i, rstart = row[j]; j < n; j++) { 79 if (row[j] != rstart) break; 80 } 81 if (j < n) ncols = j - i; 82 else ncols = n - i; 83 /* Now assemble all these values with a single function call */ 84 PetscCall(MatSetValues_MPI_Hash(A, 1, row + i, ncols, col + i, val + i, A->insertmode)); 85 i = j; 86 } 87 } 88 PetscCall(MatStashScatterEnd_Private(&A->stash)); 89 if (type != MAT_FINAL_ASSEMBLY) PetscFunctionReturn(PETSC_SUCCESS); 90 91 A->insertmode = NOT_SET_VALUES; /* this was set by the previous calls to MatSetValues() */ 92 93 PetscCall(PetscMemcpy(&A->ops, &a->cops, sizeof(*(A->ops)))); 94 A->hash_active = PETSC_FALSE; 95 96 PetscCall(MatAssemblyBegin(a->A, MAT_FINAL_ASSEMBLY)); 97 PetscCall(MatAssemblyEnd(a->A, MAT_FINAL_ASSEMBLY)); 98 PetscCall(MatAssemblyBegin(a->B, MAT_FINAL_ASSEMBLY)); 99 PetscCall(MatAssemblyEnd(a->B, MAT_FINAL_ASSEMBLY)); 100 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 101 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 102 PetscFunctionReturn(PETSC_SUCCESS); 103 } 104 105 static PetscErrorCode MatDestroy_MPI_Hash(Mat A) 106 { 107 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 108 109 PetscFunctionBegin; 110 PetscCall(MatStashDestroy_Private(&A->stash)); 111 PetscCall(MatDestroy(&a->A)); 112 PetscCall(MatDestroy(&a->B)); 113 PetscCall((*a->cops.destroy)(A)); 114 PetscFunctionReturn(PETSC_SUCCESS); 115 } 116 117 static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A) 118 { 119 PetscFunctionBegin; 120 PetscFunctionReturn(PETSC_SUCCESS); 121 } 122 123 static PetscErrorCode MatSetRandom_MPI_Hash(Mat A, PETSC_UNUSED PetscRandom r) 124 { 125 SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Must set preallocation first"); 126 } 127 128 static PetscErrorCode MatSetUp_MPI_Hash(Mat A) 129 { 130 PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data; 131 PetscMPIInt size; 132 #if !defined(TYPE_AIJ) 133 PetscInt bs; 134 #endif 135 136 PetscFunctionBegin; 137 PetscCall(PetscInfo(A, "Using hash-based MatSetValues() for MATMPISBAIJ because no preallocation provided\n")); 138 PetscCall(PetscLayoutSetUp(A->rmap)); 139 PetscCall(PetscLayoutSetUp(A->cmap)); 140 if (A->rmap->bs < 1) A->rmap->bs = 1; 141 if (A->cmap->bs < 1) A->cmap->bs = 1; 142 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size)); 143 144 #if !defined(TYPE_AIJ) 145 PetscCall(MatGetBlockSize(A, &bs)); 146 /* these values are set in MatMPISBAIJSetPreallocation() */ 147 a->bs2 = bs * bs; 148 a->mbs = A->rmap->n / bs; 149 a->nbs = A->cmap->n / bs; 150 a->Mbs = A->rmap->N / bs; 151 a->Nbs = A->cmap->N / bs; 152 153 for (PetscInt i = 0; i <= a->size; i++) a->rangebs[i] = A->rmap->range[i] / bs; 154 a->rstartbs = A->rmap->rstart / bs; 155 a->rendbs = A->rmap->rend / bs; 156 a->cstartbs = A->cmap->rstart / bs; 157 a->cendbs = A->cmap->rend / bs; 158 PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), A->rmap->bs, &A->bstash)); 159 #endif 160 161 PetscCall(MatCreate(PETSC_COMM_SELF, &a->A)); 162 PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->n, A->rmap->n, A->cmap->n)); 163 PetscCall(MatSetBlockSizesFromMats(a->A, A, A)); 164 #if defined(SUB_TYPE_CUSPARSE) 165 PetscCall(MatSetType(a->A, MATSEQAIJCUSPARSE)); 166 #else 167 PetscCall(MatSetType(a->A, PetscConcat(MATSEQ, TYPE))); 168 #endif 169 PetscCall(MatSetUp(a->A)); 170 171 PetscCall(MatCreate(PETSC_COMM_SELF, &a->B)); 172 PetscCall(MatSetSizes(a->B, A->rmap->n, size > 1 ? A->cmap->N : 0, A->rmap->n, size > 1 ? A->cmap->N : 0)); 173 PetscCall(MatSetBlockSizesFromMats(a->B, A, A)); 174 #if defined(TYPE_SBAIJ) 175 PetscCall(MatSetType(a->B, MATSEQBAIJ)); 176 #else 177 #if defined(SUB_TYPE_CUSPARSE) 178 PetscCall(MatSetType(a->B, MATSEQAIJCUSPARSE)); 179 #else 180 PetscCall(MatSetType(a->B, PetscConcat(MATSEQ, TYPE))); 181 #endif 182 #endif 183 PetscCall(MatSetUp(a->B)); 184 185 /* keep a record of the operations so they can be reset when the hash handling is complete */ 186 PetscCall(PetscMemcpy(&a->cops, &A->ops, sizeof(*(A->ops)))); 187 188 A->ops->assemblybegin = MatAssemblyBegin_MPI_Hash; 189 A->ops->assemblyend = MatAssemblyEnd_MPI_Hash; 190 A->ops->setvalues = MatSetValues_MPI_Hash; 191 A->ops->destroy = MatDestroy_MPI_Hash; 192 A->ops->zeroentries = MatZeroEntries_MPI_Hash; 193 A->ops->setrandom = MatSetRandom_MPI_Hash; 194 A->ops->setvaluesblocked = NULL; 195 196 A->preallocated = PETSC_TRUE; 197 A->hash_active = PETSC_TRUE; 198 PetscFunctionReturn(PETSC_SUCCESS); 199 } 200