1: #include <petsc/private/matimpl.h>
  2: #include <petsc/private/vecimpl.h>

  4: PetscErrorCode MatConvert_Shell(Mat oldmat, MatType newtype, MatReuse reuse, Mat *newmat)
  5: {
  6:   Mat          mat;
  7:   Vec          in, out;
  8:   PetscScalar *array;
  9:   PetscInt    *dnnz, *onnz, *dnnzu, *onnzu;
 10:   PetscInt     cst, cen, Nbs, mbs, nbs, rbs, cbs;
 11:   PetscInt     im, i, m, n, M, N, *rows, start;

 13:   PetscFunctionBegin;
 14:   PetscCall(MatGetOwnershipRange(oldmat, &start, NULL));
 15:   PetscCall(MatGetOwnershipRangeColumn(oldmat, &cst, &cen));
 16:   PetscCall(MatCreateVecs(oldmat, &in, &out));
 17:   PetscCall(MatGetLocalSize(oldmat, &m, &n));
 18:   PetscCall(MatGetSize(oldmat, &M, &N));
 19:   PetscCall(PetscMalloc1(m, &rows));
 20:   if (reuse != MAT_REUSE_MATRIX) {
 21:     PetscCall(MatCreate(PetscObjectComm((PetscObject)oldmat), &mat));
 22:     PetscCall(MatSetSizes(mat, m, n, M, N));
 23:     PetscCall(MatSetType(mat, newtype));
 24:     PetscCall(MatSetBlockSizesFromMats(mat, oldmat, oldmat));
 25:     PetscCall(MatGetBlockSizes(mat, &rbs, &cbs));
 26:     mbs = m / rbs;
 27:     nbs = n / cbs;
 28:     Nbs = N / cbs;
 29:     cst = cst / cbs;
 30:     PetscCall(PetscMalloc4(mbs, &dnnz, mbs, &onnz, mbs, &dnnzu, mbs, &onnzu));
 31:     for (i = 0; i < mbs; i++) {
 32:       dnnz[i]  = nbs;
 33:       onnz[i]  = Nbs - nbs;
 34:       dnnzu[i] = PetscMax(nbs - i, 0);
 35:       onnzu[i] = PetscMax(Nbs - (cst + nbs), 0);
 36:     }
 37:     PetscCall(MatXAIJSetPreallocation(mat, PETSC_DECIDE, dnnz, onnz, dnnzu, onnzu));
 38:     PetscCall(PetscFree4(dnnz, onnz, dnnzu, onnzu));
 39:     PetscCall(VecSetOption(in, VEC_IGNORE_OFF_PROC_ENTRIES, PETSC_TRUE));
 40:     PetscCall(MatSetUp(mat));
 41:   } else {
 42:     mat = *newmat;
 43:     PetscCall(MatZeroEntries(mat));
 44:   }
 45:   for (i = 0; i < N; i++) {
 46:     PetscInt j;

 48:     PetscCall(VecZeroEntries(in));
 49:     if (in->ops->setvalues) PetscCall(VecSetValue(in, i, 1., INSERT_VALUES));
 50:     else {
 51:       if (i >= cst && i < cen) {
 52:         PetscCall(VecGetArray(in, &array));
 53:         array[i - cst] = 1.0;
 54:         PetscCall(VecRestoreArray(in, &array));
 55:       }
 56:     }
 57:     PetscCall(VecAssemblyBegin(in));
 58:     PetscCall(VecAssemblyEnd(in));
 59:     PetscCall(MatMult(oldmat, in, out));
 60:     PetscCall(VecGetArray(out, &array));
 61:     for (j = 0, im = 0; j < m; j++) {
 62:       if (PetscAbsScalar(array[j]) == 0.0) continue;
 63:       rows[im]  = j + start;
 64:       array[im] = array[j];
 65:       im++;
 66:     }
 67:     PetscCall(MatSetValues(mat, im, rows, 1, &i, array, INSERT_VALUES));
 68:     PetscCall(VecRestoreArray(out, &array));
 69:   }
 70:   PetscCall(PetscFree(rows));
 71:   PetscCall(VecDestroy(&in));
 72:   PetscCall(VecDestroy(&out));
 73:   PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
 74:   PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));
 75:   if (reuse == MAT_INPLACE_MATRIX) {
 76:     PetscCall(MatHeaderReplace(oldmat, &mat));
 77:   } else {
 78:     *newmat = mat;
 79:   }
 80:   PetscFunctionReturn(PETSC_SUCCESS);
 81: }

 83: static PetscErrorCode MatGetDiagonal_CF(Mat A, Vec X)
 84: {
 85:   Mat B;

 87:   PetscFunctionBegin;
 88:   PetscCall(MatShellGetContext(A, &B));
 89:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
 90:   PetscCall(MatGetDiagonal(B, X));
 91:   PetscFunctionReturn(PETSC_SUCCESS);
 92: }

 94: static PetscErrorCode MatMult_CF(Mat A, Vec X, Vec Y)
 95: {
 96:   Mat B;

 98:   PetscFunctionBegin;
 99:   PetscCall(MatShellGetContext(A, &B));
100:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
101:   PetscCall(MatMult(B, X, Y));
102:   PetscFunctionReturn(PETSC_SUCCESS);
103: }

105: static PetscErrorCode MatMultTranspose_CF(Mat A, Vec X, Vec Y)
106: {
107:   Mat B;

109:   PetscFunctionBegin;
110:   PetscCall(MatShellGetContext(A, &B));
111:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
112:   PetscCall(MatMultTranspose(B, X, Y));
113:   PetscFunctionReturn(PETSC_SUCCESS);
114: }

116: static PetscErrorCode MatDestroy_CF(Mat A)
117: {
118:   Mat B;

120:   PetscFunctionBegin;
121:   PetscCall(MatShellGetContext(A, &B));
122:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
123:   PetscCall(MatDestroy(&B));
124:   PetscCall(MatShellSetContext(A, NULL));
125:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_anytype_C", NULL));
126:   PetscFunctionReturn(PETSC_SUCCESS);
127: }

129: typedef struct {
130:   void              *ctx;
131:   PetscCtxDestroyFn *destroy;
132:   PetscErrorCode (*numeric)(Mat);
133:   MatProductType ptype;
134:   Mat            Dwork;
135: } MatProductCtx_MatMatCF;

137: static PetscErrorCode MatProductDestroy_CF(PetscCtxRt data)
138: {
139:   MatProductCtx_MatMatCF *mmcfdata = *(MatProductCtx_MatMatCF **)data;

141:   PetscFunctionBegin;
142:   if (mmcfdata->destroy) PetscCall((*mmcfdata->destroy)(&mmcfdata->ctx));
143:   PetscCall(MatDestroy(&mmcfdata->Dwork));
144:   PetscCall(PetscFree(mmcfdata));
145:   PetscFunctionReturn(PETSC_SUCCESS);
146: }

148: static PetscErrorCode MatProductNumericPhase_CF(Mat A, Mat B, Mat C, void *data)
149: {
150:   MatProductCtx_MatMatCF *mmcfdata = (MatProductCtx_MatMatCF *)data;

152:   PetscFunctionBegin;
153:   PetscCheck(mmcfdata, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing data");
154:   PetscCheck(mmcfdata->numeric, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing numeric operation");
155:   /* the MATSHELL interface allows us to play with the product data */
156:   PetscCall(PetscNew(&C->product));
157:   C->product->type  = mmcfdata->ptype;
158:   C->product->data  = mmcfdata->ctx;
159:   C->product->Dwork = mmcfdata->Dwork;
160:   PetscCall(MatShellGetContext(A, &C->product->A));
161:   C->product->B = B;
162:   PetscCall((*mmcfdata->numeric)(C));
163:   PetscCall(PetscFree(C->product));
164:   PetscFunctionReturn(PETSC_SUCCESS);
165: }

167: static PetscErrorCode MatProductSymbolicPhase_CF(Mat A, Mat B, Mat C, void **data)
168: {
169:   MatProductCtx_MatMatCF *mmcfdata;

171:   PetscFunctionBegin;
172:   PetscCall(MatShellGetContext(A, &C->product->A));
173:   PetscCall(MatProductSetFromOptions(C));
174:   PetscCall(MatProductSymbolic(C));
175:   /* the MATSHELL interface does not allow non-empty product data */
176:   PetscCall(PetscNew(&mmcfdata));

178:   mmcfdata->numeric = C->ops->productnumeric;
179:   mmcfdata->ptype   = C->product->type;
180:   mmcfdata->ctx     = C->product->data;
181:   mmcfdata->destroy = C->product->destroy;
182:   mmcfdata->Dwork   = C->product->Dwork;

184:   C->product->Dwork   = NULL;
185:   C->product->data    = NULL;
186:   C->product->destroy = NULL;
187:   C->product->A       = A;

189:   *data = mmcfdata;
190:   PetscFunctionReturn(PETSC_SUCCESS);
191: }

193: /* only for A of type shell, mainly used for MatMat operations of shells with AXPYs */
194: static PetscErrorCode MatProductSetFromOptions_CF(Mat D)
195: {
196:   Mat A, B, Ain;
197:   PetscErrorCode (*Af)(Mat) = NULL;
198:   PetscBool flg;

200:   PetscFunctionBegin;
201:   MatCheckProduct(D, 1);
202:   if (D->product->type == MATPRODUCT_ABC) PetscFunctionReturn(PETSC_SUCCESS);
203:   A = D->product->A;
204:   B = D->product->B;
205:   PetscCall(MatIsShell(A, &flg));
206:   if (!flg) PetscFunctionReturn(PETSC_SUCCESS);
207:   PetscCall(PetscObjectQueryFunction((PetscObject)A, "MatProductSetFromOptions_anytype_C", &Af));
208:   if (Af == MatProductSetFromOptions_CF) PetscCall(MatShellGetContext(A, &Ain));
209:   else PetscFunctionReturn(PETSC_SUCCESS);
210:   D->product->A = Ain;
211:   PetscCall(MatProductSetFromOptions(D));
212:   D->product->A = A;
213:   if (D->ops->productsymbolic) { /* we have a symbolic match, now populate the MATSHELL operations */
214:     PetscCall(MatShellSetMatProductOperation(A, D->product->type, MatProductSymbolicPhase_CF, MatProductNumericPhase_CF, MatProductDestroy_CF, ((PetscObject)B)->type_name, NULL));
215:     PetscCall(MatProductSetFromOptions(D));
216:   }
217:   PetscFunctionReturn(PETSC_SUCCESS);
218: }

220: PetscErrorCode MatConvertFrom_Shell(Mat A, MatType newtype, MatReuse reuse, Mat *B)
221: {
222:   Mat       M;
223:   PetscBool flg;

225:   PetscFunctionBegin;
226:   PetscCall(PetscStrcmp(newtype, MATSHELL, &flg));
227:   PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only conversion to MATSHELL");
228:   if (reuse == MAT_INITIAL_MATRIX) {
229:     PetscCall(PetscObjectReference((PetscObject)A));
230:     PetscCall(MatCreateShell(PetscObjectComm((PetscObject)A), A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N, A, &M));
231:     PetscCall(MatSetBlockSizesFromMats(M, A, A));
232:     PetscCall(MatShellSetOperation(M, MATOP_MULT, (PetscErrorCodeFn *)MatMult_CF));
233:     PetscCall(MatShellSetOperation(M, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_CF));
234:     PetscCall(MatShellSetOperation(M, MATOP_GET_DIAGONAL, (PetscErrorCodeFn *)MatGetDiagonal_CF));
235:     PetscCall(MatShellSetOperation(M, MATOP_DESTROY, (PetscErrorCodeFn *)MatDestroy_CF));
236:     PetscCall(PetscObjectComposeFunction((PetscObject)M, "MatProductSetFromOptions_anytype_C", MatProductSetFromOptions_CF));
237:     PetscCall(PetscFree(M->defaultvectype));
238:     PetscCall(PetscStrallocpy(A->defaultvectype, &M->defaultvectype));
239: #if defined(PETSC_HAVE_DEVICE)
240:     PetscCall(MatBindToCPU(M, A->boundtocpu));
241: #endif
242:     *B = M;
243:   } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Not implemented");
244:   PetscFunctionReturn(PETSC_SUCCESS);
245: }