Actual source code: ex304.c
1: static char help[] = "Test matmat products with matdiagonal on gpus \n\n";
3: // Contributed by: Steven Dargaville
5: #include <petscmat.h>
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;
15: PetscFunctionBeginUser;
16: PetscCall(PetscInitialize(&argc, &args, NULL, help));
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));
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));
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));
39: // ~~~~~~~~~~~~~
40: // Do an initial matmatmult
41: // A * B_aij_diag
42: // and then
43: // A * B_diag but just using MatDiagonalScale
44: // ~~~~~~~~~~~~~
46: // aij * aij
47: PetscCall(MatMatMult(A, B_aij_diag, MAT_INITIAL_MATRIX, 1.5, &result));
48: // PetscCall(MatView(result, PETSC_VIEWER_STDOUT_WORLD));
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));
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");
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));
68: // aij * aij
69: PetscCall(MatMatMult(A, B_aij_diag, MAT_REUSE_MATRIX, 1.5, &result));
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));
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");
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));
87: PetscCall(PetscFinalize());
88: return 0;
89: }
91: /*TEST
92: test:
93: requires: kokkos_kernels
94: args: -mat_type aijkokkos
95: output_file: output/empty.out
96: TEST*/