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