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