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