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