1 static char help[] = "Test matmat products with matdiagonal on gpus \n\n";
2
3 // Contributed by: Steven Dargaville
4
5 #include <petscmat.h>
6
main(int argc,char ** args)7 int main(int argc, char **args)
8 {
9 const PetscInt inds[] = {0, 1};
10 PetscScalar avals[] = {2, 3, 5, 7};
11 Mat A, B_diag, B_aij_diag, result, result_diag;
12 Vec diag;
13 PetscBool equal = PETSC_FALSE;
14
15 PetscFunctionBeginUser;
16 PetscCall(PetscInitialize(&argc, &args, NULL, help));
17
18 // Create matrix to start
19 PetscCall(MatCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, 2, 2, 2, 2, &A));
20 PetscCall(MatSetUp(A));
21 PetscCall(MatSetValues(A, 2, inds, 2, inds, avals, INSERT_VALUES));
22 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
23 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
24
25 // Create a matdiagonal matrix
26 // Will be the matching vec type as A
27 PetscCall(MatCreateVecs(A, &diag, NULL));
28 PetscCall(VecSet(diag, 2.0));
29 PetscCall(MatCreateDiagonal(diag, &B_diag));
30
31 // Create the same matrix as the matdiagonal but in aij format
32 PetscCall(MatCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, 2, 2, 2, 2, &B_aij_diag));
33 PetscCall(MatSetUp(B_aij_diag));
34 PetscCall(MatDiagonalSet(B_aij_diag, diag, INSERT_VALUES));
35 PetscCall(MatAssemblyBegin(B_aij_diag, MAT_FINAL_ASSEMBLY));
36 PetscCall(MatAssemblyEnd(B_aij_diag, MAT_FINAL_ASSEMBLY));
37 PetscCall(VecDestroy(&diag));
38
39 // ~~~~~~~~~~~~~
40 // Do an initial matmatmult
41 // A * B_aij_diag
42 // and then
43 // A * B_diag but just using MatDiagonalScale
44 // ~~~~~~~~~~~~~
45
46 // aij * aij
47 PetscCall(MatMatMult(A, B_aij_diag, MAT_INITIAL_MATRIX, 1.5, &result));
48 // PetscCall(MatView(result, PETSC_VIEWER_STDOUT_WORLD));
49
50 // aij * diagonal
51 PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &result_diag));
52 PetscCall(MatDiagonalGetDiagonal(B_diag, &diag));
53 PetscCall(MatDiagonalScale(result_diag, NULL, diag));
54 PetscCall(MatDiagonalRestoreDiagonal(B_diag, &diag));
55 // PetscCall(MatView(result_diag, PETSC_VIEWER_STDOUT_WORLD));
56
57 PetscCall(MatEqual(result, result_diag, &equal));
58 PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatMatMult and MatDiagonalScale do not give the same result");
59
60 // ~~~~~~~~~~~~~
61 // Now let's modify the diagonal and do it again with "reuse"
62 // ~~~~~~~~~~~~~
63 PetscCall(MatDiagonalGetDiagonal(B_diag, &diag));
64 PetscCall(VecSet(diag, 3.0));
65 PetscCall(MatDiagonalSet(B_aij_diag, diag, INSERT_VALUES));
66 PetscCall(MatDiagonalRestoreDiagonal(B_diag, &diag));
67
68 // aij * aij
69 PetscCall(MatMatMult(A, B_aij_diag, MAT_REUSE_MATRIX, 1.5, &result));
70
71 // aij * diagonal
72 PetscCall(MatCopy(A, result_diag, SAME_NONZERO_PATTERN));
73 PetscCall(MatDiagonalGetDiagonal(B_diag, &diag));
74 PetscCall(MatDiagonalScale(result_diag, NULL, diag));
75 PetscCall(MatDiagonalRestoreDiagonal(B_diag, &diag));
76
77 PetscCall(MatEqual(result, result_diag, &equal));
78 PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatMatMult and MatDiagonalScale do not give the same result");
79
80 PetscCall(MatDestroy(&A));
81 PetscCall(MatDestroy(&B_diag));
82 PetscCall(MatDestroy(&B_aij_diag));
83 PetscCall(MatDestroy(&result));
84 PetscCall(MatDestroy(&result_diag));
85 PetscCall(VecDestroy(&diag));
86
87 PetscCall(PetscFinalize());
88 return 0;
89 }
90
91 /*TEST
92 test:
93 requires: kokkos_kernels
94 args: -mat_type aijkokkos
95 output_file: output/empty.out
96 TEST*/
97