1 2 static char help[] = "Shows how to add a new MatOperation to AIJ MatType\n\n"; 3 4 #include <petscmat.h> 5 #include <petscblaslapack.h> 6 7 static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA, PetscScalar alpha) { 8 PetscFunctionBegin; 9 PetscCall(MatScale(inA, alpha)); 10 PetscFunctionReturn(0); 11 } 12 13 extern PetscErrorCode MatScaleUserImpl(Mat, PetscScalar); 14 15 static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A, PetscScalar aa) { 16 Mat AA, AB; 17 18 PetscFunctionBegin; 19 PetscCall(MatMPIAIJGetSeqAIJ(A, &AA, &AB, NULL)); 20 PetscCall(MatScaleUserImpl(AA, aa)); 21 PetscCall(MatScaleUserImpl(AB, aa)); 22 PetscFunctionReturn(0); 23 } 24 25 /* This routine registers MatScaleUserImpl_SeqAIJ() and 26 MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl() 27 functionality for SeqAIJ and MPIAIJ matrix-types */ 28 PetscErrorCode RegisterMatScaleUserImpl(Mat mat) { 29 PetscMPIInt size; 30 31 PetscFunctionBegin; 32 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size)); 33 if (size == 1) { /* SeqAIJ Matrix */ 34 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ)); 35 } else { /* MPIAIJ Matrix */ 36 Mat AA, AB; 37 PetscCall(MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL)); 38 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_MPIAIJ)); 39 PetscCall(PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ)); 40 PetscCall(PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ)); 41 } 42 PetscFunctionReturn(0); 43 } 44 45 /* This routine deregisters MatScaleUserImpl_SeqAIJ() and 46 MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl() 47 functionality for SeqAIJ and MPIAIJ matrix-types */ 48 PetscErrorCode DeRegisterMatScaleUserImpl(Mat mat) { 49 PetscMPIInt size; 50 51 PetscFunctionBegin; 52 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size)); 53 if (size == 1) { /* SeqAIJ Matrix */ 54 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL)); 55 } else { /* MPIAIJ Matrix */ 56 Mat AA, AB; 57 PetscCall(MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL)); 58 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL)); 59 PetscCall(PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", NULL)); 60 PetscCall(PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", NULL)); 61 } 62 PetscFunctionReturn(0); 63 } 64 65 /* this routines queries the already registered MatScaleUserImp_XXX 66 implementations for the given matrix, and calls the correct 67 routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets 68 called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets 69 called */ 70 PetscErrorCode MatScaleUserImpl(Mat mat, PetscScalar a) { 71 PetscErrorCode (*f)(Mat, PetscScalar); 72 73 PetscFunctionBegin; 74 PetscCall(PetscObjectQueryFunction((PetscObject)mat, "MatScaleUserImpl_C", &f)); 75 if (f) PetscCall((*f)(mat, a)); 76 PetscFunctionReturn(0); 77 } 78 79 /* Main user code that uses MatScaleUserImpl() */ 80 81 int main(int argc, char **args) { 82 Mat mat; 83 PetscInt i, j, m = 2, n, Ii, J; 84 PetscScalar v, none = -1.0; 85 PetscMPIInt rank, size; 86 87 PetscFunctionBeginUser; 88 PetscCall(PetscInitialize(&argc, &args, (char *)0, help)); 89 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank)); 90 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size)); 91 n = 2 * size; 92 93 /* create the matrix */ 94 PetscCall(MatCreate(PETSC_COMM_WORLD, &mat)); 95 PetscCall(MatSetSizes(mat, PETSC_DECIDE, PETSC_DECIDE, m * n, m * n)); 96 PetscCall(MatSetType(mat, MATAIJ)); 97 PetscCall(MatSetUp(mat)); 98 99 /* register user defined MatScaleUser() operation for both SeqAIJ 100 and MPIAIJ types */ 101 PetscCall(RegisterMatScaleUserImpl(mat)); 102 103 /* assemble the matrix */ 104 for (i = 0; i < m; i++) { 105 for (j = 2 * rank; j < 2 * rank + 2; j++) { 106 v = -1.0; 107 Ii = j + n * i; 108 if (i > 0) { 109 J = Ii - n; 110 PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES)); 111 } 112 if (i < m - 1) { 113 J = Ii + n; 114 PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES)); 115 } 116 if (j > 0) { 117 J = Ii - 1; 118 PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES)); 119 } 120 if (j < n - 1) { 121 J = Ii + 1; 122 PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES)); 123 } 124 v = 4.0; 125 PetscCall(MatSetValues(mat, 1, &Ii, 1, &Ii, &v, INSERT_VALUES)); 126 } 127 } 128 PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY)); 129 PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY)); 130 131 /* check the matrix before and after scaling by -1.0 */ 132 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Matrix _before_ MatScaleUserImpl() operation\n")); 133 PetscCall(MatView(mat, PETSC_VIEWER_STDOUT_WORLD)); 134 PetscCall(MatScaleUserImpl(mat, none)); 135 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Matrix _after_ MatScaleUserImpl() operation\n")); 136 PetscCall(MatView(mat, PETSC_VIEWER_STDOUT_WORLD)); 137 138 /* deregister user defined MatScaleUser() operation for both SeqAIJ 139 and MPIAIJ types */ 140 PetscCall(DeRegisterMatScaleUserImpl(mat)); 141 PetscCall(MatDestroy(&mat)); 142 PetscCall(PetscFinalize()); 143 return 0; 144 } 145 146 /*TEST 147 148 test: 149 150 test: 151 suffix: 2 152 nsize: 2 153 154 TEST*/ 155