1: static char help[] = "Testing MatCreateMPIAIJWithSeqAIJ().\n\n";

  3: #include <petscmat.h>

  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;

 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");

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

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

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

 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

 51:   // 2) Test MatCreateMPIAIJWithSeqAIJ with non-compressed off-diag.
 52:   PetscCall(MatConvert(A, mat_type, MAT_INITIAL_MATRIX, &A2));

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

 67:   PetscCall(MatRestoreRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nd, (const PetscInt **)&bi, (const PetscInt **)&bj, &done));
 68:   PetscCall(MatSeqAIJRestoreArrayRead(B, &ba));

 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()");

 75:   PetscCall(MatDestroy(&mat));
 76:   PetscCall(MatDestroy(&mat2));
 77:   PetscCall(PetscFinalize());
 78:   return 0;
 79: }

 81: /*TEST

 83:   testset:
 84:     nsize: 2
 85:     output_file: output/empty.out

 87:     test:
 88:       suffix: 1
 89:       args: -mat_type aij

 91:     test:
 92:       suffix: 2
 93:       args: -mat_type aijkokkos
 94:       requires: kokkos_kernels

 96:     test:
 97:       suffix: 3
 98:       args: -mat_type aijcusparse
 99:       requires: cuda

101:     test:
102:       suffix: 4
103:       args: -mat_type aijhipsparse
104:       requires: hip

106: TEST*/