TpetraCrsColorer.cpp

Go to the documentation of this file.

#include "Tpetra_Core.hpp"
#include "Kokkos_Random.hpp"
#include "Zoltan2_TestHelpers.hpp"
#include "Zoltan2_TpetraCrsColorer.hpp"

// Class to test the Colorer utility
class ColorerTest {
public:
  using map_t = Tpetra::Map<>;
  using gno_t = typename map_t::global_ordinal_type;
  using graph_t = Tpetra::CrsGraph<>;
  using matrix_t = Tpetra::CrsMatrix<zscalar_t>;
  using multivector_t = Tpetra::MultiVector<zscalar_t>;
  using execution_space_t = typename matrix_t::device_type::execution_space;

  // Construct the test:
  //   Read or generate a matrix (JBlock) with default range and domain maps
  //   Construct identical matrix (JCyclic) with cyclic range and domain maps

  ColorerTest(Teuchos::RCP<const Teuchos::Comm<int> > &comm,
              int narg, char**arg)
  : symmetric(false)
  {
    int me = comm->getRank();
    int np = comm->getSize();

    // Process command line arguments
    bool distributeInput = true;
    size_t xdim = 10, ydim = 11, zdim = 12;

    Teuchos::CommandLineProcessor cmdp(false, false);
    cmdp.setOption("file", &matrixFileName,
                   "Name of the Matrix Market file to use");
    cmdp.setOption("xdim", &xdim,
                   "Number of nodes in x-direction for generated matrix");
    cmdp.setOption("ydim", &ydim,
                   "Number of nodes in y-direction for generated matrix");
    cmdp.setOption("zdim", &zdim,
                   "Number of nodes in z-direction for generated matrix");
    cmdp.setOption("distribute", "no-distribute", &distributeInput,
                   "Should Zoltan2 distribute the matrix as it is read?");
    cmdp.setOption("symmetric", "non-symmetric", &symmetric,
                   "Is the matrix symmetric?");
    cmdp.parse(narg, arg);

    // Get and store a matrix
    if (matrixFileName != "") {
      // Read from a file
      UserInputForTests uinput(".", matrixFileName, comm, true, distributeInput);
      JBlock = uinput.getUITpetraCrsMatrix();
    }
    else {
      // Generate a matrix
      UserInputForTests uinput(xdim, ydim, zdim, string("Laplace3D"), comm,
                               true, distributeInput);
      JBlock = uinput.getUITpetraCrsMatrix();
    }

    // Build same matrix with cyclic domain and range maps
    // To make range and domain maps differ for square matrices,
    // keep some processors empty in the cyclic maps

    size_t nIndices = std::max(JBlock->getGlobalNumCols(),
                               JBlock->getGlobalNumRows());
    Teuchos::Array<gno_t> indices(nIndices);

    Teuchos::RCP<const map_t> vMapCyclic =
                 getCyclicMap(JBlock->getGlobalNumCols(), indices, np-1, comm);
    Teuchos::RCP<const map_t> wMapCyclic =
                 getCyclicMap(JBlock->getGlobalNumRows(), indices, np-2, comm);

    // Fill JBlock with random numbers for a better test.
    JBlock->resumeFill();

    using IST = typename Kokkos::ArithTraits<zscalar_t>::val_type;
    using pool_type =
          Kokkos::Random_XorShift64_Pool<execution_space_t>;
    pool_type rand_pool(static_cast<uint64_t>(me));

    Kokkos::fill_random(JBlock->getLocalMatrixDevice().values, rand_pool,
                        static_cast<IST>(1.), static_cast<IST>(9999.));
    JBlock->fillComplete();


    // Make JCyclic:  same matrix with different Domain and Range maps
    RCP<const graph_t> block_graph = JBlock->getCrsGraph();
    RCP<graph_t> cyclic_graph = rcp(new graph_t(*block_graph));
    cyclic_graph->resumeFill();
    cyclic_graph->fillComplete(vMapCyclic, wMapCyclic);
    JCyclic = rcp(new matrix_t(cyclic_graph));
    JCyclic->resumeFill();
    TEUCHOS_ASSERT(block_graph->getLocalNumRows() == cyclic_graph->getLocalNumRows());
    {
      auto val_s = JBlock->getLocalMatrixHost().values;
      auto val_d = JCyclic->getLocalMatrixHost().values;
      TEUCHOS_ASSERT(val_s.extent(0) == val_d.extent(0));
      Kokkos::deep_copy(val_d, val_s);
    }
    JCyclic->fillComplete();
  }

  bool run(const char *testname, Teuchos::ParameterList &params) {

    bool ok = true;

    params.set("symmetric", symmetric);
    params.set("library", "zoltan");

    // test with default maps
    ok = buildAndCheckSeedMatrix(testname, params, true);

    // test with cyclic maps
    ok &= buildAndCheckSeedMatrix(testname, params, false);
    // if (matrixFileName != "west0067") {

      params.set("library", "zoltan2");
      // test with default maps
      ok = buildAndCheckSeedMatrix(testname, params, true);

      // test with cyclic maps
      ok &= buildAndCheckSeedMatrix(testname, params, false);
    // }

    return ok;
  }

  bool buildAndCheckSeedMatrix(
    const char *testname,
    Teuchos::ParameterList &params,
    const bool useBlock
  )
  {
    int ierr = 0;

    // Pick matrix depending on useBlock flag
    Teuchos::RCP<matrix_t> J = (useBlock ? JBlock : JCyclic);
    int me = J->getRowMap()->getComm()->getRank();

    std::cout << params.get("library", "zoltan2") << " Running " << testname << " with "
              << (useBlock ? "Block maps" : "Cyclic maps")
              << std::endl;

    // Create a colorer
    Zoltan2::TpetraCrsColorer<matrix_t> colorer(J);
    colorer.computeColoring(params);

    // Check coloring
    if (!colorer.checkColoring()) {
      std::cout << testname << " with "
                << (useBlock ? "Block maps" : "Cyclic maps")
                << " FAILED: invalid coloring returned"
                << std::endl;
      return false;
    }

    // Compute seed matrix V -- the application wants this matrix
    // Dense matrix of 0/1 indicating the compression via coloring

    const int numColors = colorer.getNumColors();

    // Compute the seed matrix; applications want this seed matrix

    multivector_t V(J->getDomainMap(), numColors);
    colorer.computeSeedMatrix(V);

    // To test the result...
    // Compute the compressed matrix W
    multivector_t W(J->getRangeMap(), numColors);

    J->apply(V, W);

    // Reconstruct matrix from compression vector
    Teuchos::RCP<matrix_t> Jp = rcp(new matrix_t(*J, Teuchos::Copy));
    Jp->setAllToScalar(static_cast<zscalar_t>(-1.));

    colorer.reconstructMatrix(W, *Jp);

    // Check that values of J = values of Jp
    auto J_local_matrix = J->getLocalMatrixDevice();
    auto Jp_local_matrix = Jp->getLocalMatrixDevice();
    const size_t num_local_nz = J->getLocalNumEntries();

    Kokkos::parallel_reduce(
      "TpetraCrsColorer::testReconstructedMatrix()",
      Kokkos::RangePolicy<execution_space_t>(0, num_local_nz),
      KOKKOS_LAMBDA(const size_t nz, int &errorcnt) {
        if (J_local_matrix.values(nz) != Jp_local_matrix.values(nz)) {
          Kokkos::printf("Error in nonzero comparison %zu:  %g != %g",
                  nz, J_local_matrix.values(nz), Jp_local_matrix.values(nz));
          errorcnt++;
        }
      },
      ierr);


    if (ierr > 0) {
      std::cout << testname << " FAILED on rank " << me << " with "
                << (useBlock ? "Block maps" : "Cyclic maps")
                << std::endl;
      params.print();
    }

    return (ierr == 0);
  }

private:

  // Return a map that is cyclic (like dealing rows to processors)
  Teuchos::RCP<const map_t> getCyclicMap(
    size_t nIndices,
    Teuchos::Array<gno_t> &indices,
    int mapNumProc,
    const Teuchos::RCP<const Teuchos::Comm<int> > &comm)
  {
    size_t cnt = 0;
    int me = comm->getRank();
    int np = comm->getSize();
    if (mapNumProc > np) mapNumProc = np; // corner case: bad input
    if (mapNumProc <= 0) mapNumProc = 1;  // corner case: np is too small

    for (size_t i = 0; i < nIndices; i++)
      if (me == int(i % np)) indices[cnt++] = i;

    Tpetra::global_size_t dummy =
            Teuchos::OrdinalTraits<Tpetra::global_size_t>::invalid();

    return rcp(new map_t(dummy, indices(0,cnt), 0, comm));
  }

  // Input matrix -- built in Constructor
  bool symmetric;               // User can specify whether matrix is symmetric
  Teuchos::RCP<matrix_t> JBlock;   // has Trilinos default domain and range maps
  Teuchos::RCP<matrix_t> JCyclic;  // has cyclic domain and range maps
  std::string matrixFileName;
};


int main(int narg, char **arg)
{
  Tpetra::ScopeGuard scope(&narg, &arg);
  Teuchos::RCP<const Teuchos::Comm<int> > comm = Tpetra::getDefaultComm();
  bool ok = true;
  int ierr = 0;

  ColorerTest testColorer(comm, narg, arg);

  // Set parameters and compute coloring
  {
    Teuchos::ParameterList coloring_params;
    std::string matrixType = "Jacobian";
    bool symmetrize = true;  // Zoltan's TRANSPOSE symmetrization, if needed

    coloring_params.set("MatrixType", matrixType);
    coloring_params.set("symmetrize", symmetrize);

    ok = testColorer.run("Test One", coloring_params);
    if (!ok) ierr++;
  }

  {
    Teuchos::ParameterList coloring_params;
    std::string matrixType = "Jacobian";
    bool symmetrize = false;  // colorer's BIPARTITE symmetrization, if needed

    coloring_params.set("MatrixType", matrixType);
    coloring_params.set("symmetrize", symmetrize);

    ok = testColorer.run("Test Two", coloring_params);
    if (!ok) ierr++;
  }

  {
    Teuchos::ParameterList coloring_params;
    std::string matrixType = "Jacobian";

    coloring_params.set("MatrixType", matrixType);

    ok = testColorer.run("Test Three", coloring_params);
    if (!ok) ierr++;
  }

  int gerr;
  Teuchos::reduceAll<int, int>(*comm, Teuchos::REDUCE_SUM, 1, &ierr, &gerr);
  if (comm->getRank() == 0) {
    if (gerr == 0)
      std::cout << "TEST PASSED" << std::endl;
    else
      std::cout << "TEST FAILED" << std::endl;
  }

//Through cmake...
//Test cases -- UserInputForTests can generate Galeri or read files:
//-  tri-diagonal matrix -- can check the number of colors
//-  galeri matrix
//-  read from file:  symmetric matrix and non-symmetric matrix

//Through code ...
//Test with fitted and non-fitted maps
//Call regular and fitted versions of functions

//Through code ...
//Test both with and without Symmetrize --
//test both to exercise both sets of callbacks in Zoltan
// --matrixType = Jacobian/Hessian
// --symmetric, --no-symmetric
// --symmetrize, --no-symmetrize

//Through cmake
//Test both with and without distributeInput
// --distribute, --no-distribute

}