1 static char help[] = "Testing MatCreateMPIAIJWithSeqAIJ().\n\n"; 2 3 #include <petscmat.h> 4 5 int main(int argc, char **argv) 6 { 7 Mat mat, mat2; 8 Mat A, B; // diag, offdiag of mat 9 Mat A2, B2; // diag, offdiag of mat2 10 PetscInt M, N, m = 5, n = 6, d_nz = 3, o_nz = 4; 11 PetscInt *bi, *bj, nd; 12 const PetscScalar *ba; 13 const PetscInt *garray; 14 PetscInt *garray_h; 15 PetscBool equal, done; 16 PetscMPIInt size; 17 char mat_type[PETSC_MAX_PATH_LEN]; 18 PetscBool flg; 19 20 PetscFunctionBeginUser; 21 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 22 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size)); 23 PetscCheck(size > 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Must run with 2 or more processes"); 24 25 // Create a MATMPIAIJ matrix for checking against 26 PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, m, n, PETSC_DECIDE, PETSC_DECIDE, d_nz, NULL, o_nz, NULL, &mat)); 27 PetscCall(MatSetRandom(mat, NULL)); 28 PetscCall(MatGetSize(mat, &M, &N)); 29 PetscCall(MatGetLocalSize(mat, &m, &n)); 30 PetscCall(MatMPIAIJGetSeqAIJ(mat, &A, &B, &garray)); 31 32 PetscCall(PetscOptionsGetString(NULL, NULL, "-mat_type", mat_type, sizeof(mat_type), &flg)); 33 if (!flg) PetscCall(PetscStrncpy(mat_type, MATSEQAIJ, sizeof(mat_type))); // Default to MATAIJ 34 PetscCall(MatConvert(A, mat_type, MAT_INITIAL_MATRIX, &A2)); // Copy A, B to A2, B2 but in the given mat_type 35 PetscCall(MatConvert(B, mat_type, MAT_INITIAL_MATRIX, &B2)); 36 37 // The garray passed in has to be on the host, but it can be created on device and copied to the host. 38 // We're just going to copy the existing host values here. 39 PetscInt nonlocalCols; 40 PetscCall(MatGetLocalSize(B, NULL, &nonlocalCols)); 41 PetscCall(PetscMalloc1(nonlocalCols, &garray_h)); 42 for (int i = 0; i < nonlocalCols; i++) garray_h[i] = garray[i]; 43 44 // 1) Test MatCreateMPIAIJWithSeqAIJ with compressed off-diag. 45 PetscCall(MatCreateMPIAIJWithSeqAIJ(PETSC_COMM_WORLD, M, N, A2, B2, garray_h, &mat2)); 46 PetscCall(MatEqual(mat, mat2, &equal)); // might directly compare value arrays 47 if (equal) PetscCall(MatMultEqual(mat, mat2, 2, &equal)); // compare via MatMult() 48 PetscCheck(equal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Likely a bug in MatCreateMPIAIJWithSeqAIJ()"); 49 PetscCall(MatDestroy(&mat2)); // A2 and B2 are also destroyed 50 51 // 2) Test MatCreateMPIAIJWithSeqAIJ with non-compressed off-diag. 52 PetscCall(MatConvert(A, mat_type, MAT_INITIAL_MATRIX, &A2)); 53 54 // Create a version of B2 of size N with global indices 55 PetscCall(MatGetRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&bi, (const PetscInt **)&bj, &done)); 56 PetscCall(MatSeqAIJGetArrayRead(B, &ba)); 57 PetscCall(MatCreate(PETSC_COMM_SELF, &B2)); 58 PetscCall(MatSetSizes(B2, m, N, m, N)); 59 PetscCall(MatSetType(B2, mat_type)); 60 PetscCall(MatSeqAIJSetPreallocation(B2, o_nz, NULL)); 61 for (int i = 0; i < m; i++) { // Fill B2 with values from B 62 for (int j = 0; j < bi[i + 1] - bi[i]; j++) PetscCall(MatSetValue(B2, i, garray[bj[bi[i] + j]], ba[bi[i] + j], INSERT_VALUES)); 63 } 64 PetscCall(MatAssemblyBegin(B2, MAT_FINAL_ASSEMBLY)); 65 PetscCall(MatAssemblyEnd(B2, MAT_FINAL_ASSEMBLY)); 66 67 PetscCall(MatRestoreRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&bi, (const PetscInt **)&bj, &done)); 68 PetscCall(MatSeqAIJRestoreArrayRead(B, &ba)); 69 70 PetscCall(MatCreateMPIAIJWithSeqAIJ(PETSC_COMM_WORLD, M, N, A2, B2, NULL, &mat2)); 71 PetscCall(MatEqual(mat, mat2, &equal)); 72 if (equal) PetscCall(MatMultEqual(mat, mat2, 2, &equal)); 73 PetscCheck(equal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Likely a bug in MatCreateMPIAIJWithSeqAIJ()"); 74 75 PetscCall(MatDestroy(&mat)); 76 PetscCall(MatDestroy(&mat2)); 77 PetscCall(PetscFinalize()); 78 return 0; 79 } 80 81 /*TEST 82 83 testset: 84 nsize: 2 85 output_file: output/empty.out 86 87 test: 88 suffix: 1 89 args: -mat_type aij 90 91 test: 92 suffix: 2 93 args: -mat_type aijkokkos 94 requires: kokkos_kernels 95 96 test: 97 suffix: 3 98 args: -mat_type aijcusparse 99 requires: cuda 100 101 test: 102 suffix: 4 103 args: -mat_type aijhipsparse 104 requires: hip 105 106 TEST*/ 107