1: static char help[] = "Tests MatGetMultPetscSF() for the parallel AIJ, BAIJ, SBAIJ, dense, and SELL matrix types.\n\n";

  3: #include <petscmat.h>
  4: #include <petscsf.h>

  6: int main(int argc, char **argv)
  7: {
  8:   Mat       A;
  9:   Vec       x, y;
 10:   PetscSF   sf;
 11:   PetscBool isdense;
 12:   PetscInt  N = 12, rstart, rend, cstart, cend, nroots, nleaves, i, ok = 1;
 13:   PetscInt *rootdata = NULL, *leafdata = NULL;

 15:   PetscFunctionBeginUser;
 16:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 17:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &N, NULL));

 19:   /* Build a symmetric tridiagonal matrix; the type is selected with -mat_type. */
 20:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 21:   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, N, N));
 22:   PetscCall(MatSetType(A, MATMPIAIJ));
 23:   PetscCall(MatSetFromOptions(A));
 24:   PetscCall(MatSetUp(A));
 25:   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
 26:   for (i = rstart; i < rend; i++) {
 27:     PetscScalar diag = 2.0, offd = -1.0;

 29:     PetscCall(MatSetValue(A, i, i, diag, INSERT_VALUES));
 30:     if (i + 1 < N) PetscCall(MatSetValue(A, i, i + 1, offd, INSERT_VALUES));
 31:     if (i - 1 >= 0) PetscCall(MatSetValue(A, i, i - 1, offd, INSERT_VALUES));
 32:   }
 33:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
 34:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));

 36:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSE, &isdense));
 37:   if (isdense) {
 38:     /* Dense builds its matrix-multiply PetscSF lazily on the first MatMult(), so trigger it first.
 39:        The SF is an allgather: every rank gathers all N global columns. */
 40:     PetscCall(MatCreateVecs(A, &x, &y));
 41:     PetscCall(VecSet(x, 1.0));
 42:     PetscCall(MatMult(A, x, y));
 43:     PetscCall(MatGetMultPetscSF(A, &sf));
 44:     PetscCheck(sf, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatGetMultPetscSF() returned NULL on a parallel dense matrix after MatMult()");
 45:     PetscCall(PetscSFGetGraph(sf, &nroots, &nleaves, NULL, NULL));
 46:     if (nleaves != N) ok = 0;
 47:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &ok, 1, MPIU_INT, MPI_LAND, PETSC_COMM_WORLD));
 48:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "MatGetMultPetscSF() OK (PetscSF valid, allgather sees all columns: %s)\n", ok ? "yes" : "no"));
 49:     PetscCall(VecDestroy(&x));
 50:     PetscCall(VecDestroy(&y));
 51:   } else {
 52:     PetscCall(MatGetMultPetscSF(A, &sf));
 53:     PetscCheck(sf, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatGetMultPetscSF() returned NULL on a parallel matrix");

 55:     /* Broadcast each owned global column index to the off-process columns that couple
 56:        to local rows; every gathered value must lie outside the local column ownership range. */
 57:     PetscCall(MatGetOwnershipRangeColumn(A, &cstart, &cend));
 58:     PetscCall(PetscSFGetGraph(sf, &nroots, &nleaves, NULL, NULL));
 59:     PetscCall(PetscMalloc2(nroots, &rootdata, nleaves, &leafdata));
 60:     for (i = 0; i < nroots; i++) rootdata[i] = cstart + i;
 61:     PetscCall(PetscSFBcastBegin(sf, MPIU_INT, rootdata, leafdata, MPI_REPLACE));
 62:     PetscCall(PetscSFBcastEnd(sf, MPIU_INT, rootdata, leafdata, MPI_REPLACE));
 63:     for (i = 0; i < nleaves; i++)
 64:       if (leafdata[i] < 0 || leafdata[i] >= N || (leafdata[i] >= cstart && leafdata[i] < cend)) ok = 0;
 65:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &ok, 1, MPIU_INT, MPI_LAND, PETSC_COMM_WORLD));
 66:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "MatGetMultPetscSF() OK (PetscSF valid, all gathered columns off-process: %s)\n", ok ? "yes" : "no"));
 67:     PetscCall(PetscFree2(rootdata, leafdata));
 68:   }
 69:   PetscCall(MatDestroy(&A));
 70:   PetscCall(PetscFinalize());
 71:   return 0;
 72: }

 74: /*TEST

 76:    test:
 77:       suffix: aij
 78:       nsize: {{2 3}}
 79:       args: -mat_type mpiaij
 80:       output_file: output/ex308.out

 82:    test:
 83:       suffix: baij
 84:       nsize: {{2 3}}
 85:       args: -mat_type mpibaij
 86:       output_file: output/ex308.out

 88:    test:
 89:       suffix: sbaij
 90:       nsize: {{2 3}}
 91:       args: -mat_type mpisbaij
 92:       output_file: output/ex308.out

 94:    test:
 95:       suffix: kokkos
 96:       requires: kokkos_kernels
 97:       nsize: {{2 3}}
 98:       args: -mat_type mpiaijkokkos
 99:       output_file: output/ex308.out

101:    test:
102:       suffix: dense
103:       nsize: {{2 3}}
104:       args: -mat_type mpidense
105:       output_file: output/ex308_dense.out

107:    test:
108:       suffix: sell
109:       nsize: {{2 3}}
110:       args: -mat_type mpisell
111:       output_file: output/ex308.out

113: TEST*/