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

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

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

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

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

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

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

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

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

130: typedef struct {
131:   void *userdata;
132:   PetscErrorCode (*userdestroy)(void *);
133:   PetscErrorCode (*numeric)(Mat);
134:   MatProductType ptype;
135:   Mat            Dwork;
136: } MatMatCF;

138: static PetscErrorCode MatProductDestroy_CF(void *data)
139: {
140:   MatMatCF *mmcfdata = (MatMatCF *)data;

142:   PetscFunctionBegin;
143:   if (mmcfdata->userdestroy) PetscCall((*mmcfdata->userdestroy)(mmcfdata->userdata));
144:   PetscCall(MatDestroy(&mmcfdata->Dwork));
145:   PetscCall(PetscFree(mmcfdata));
146:   PetscFunctionReturn(PETSC_SUCCESS);
147: }

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

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

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

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

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

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

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

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

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

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

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