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 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