TpetraRowGraphInputKokkos.cpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//   Zoltan2: A package of combinatorial algorithms for scientific computing
//
// Copyright 2012 NTESS and the Zoltan2 contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
//
// Basic testing of Zoltan2::TpetraRowGraphAdapter
#include <Zoltan2_InputTraits.hpp>
#include <Zoltan2_TestHelpers.hpp>
#include <Zoltan2_TpetraCrsGraphAdapter.hpp>
#include <Zoltan2_TpetraRowGraphAdapter.hpp>
 
#include <Teuchos_Comm.hpp>
#include <Teuchos_CommHelpers.hpp>
#include <Teuchos_DefaultComm.hpp>
#include <Teuchos_RCP.hpp>
#include <cstdlib>
#include <stdexcept>
 
using Teuchos::Comm;
using Teuchos::RCP;
using Teuchos::rcp;
using Teuchos::rcp_const_cast;
using Teuchos::rcp_dynamic_cast;
 
using ztcrsgraph_t = Tpetra::CrsGraph<zlno_t, zgno_t, znode_t>;
using ztrowgraph_t = Tpetra::RowGraph<zlno_t, zgno_t, znode_t>;
using node_t = typename Zoltan2::InputTraits<ztrowgraph_t>::node_t;
using device_t = typename node_t::device_type;
using rowAdapter_t = Zoltan2::TpetraRowGraphAdapter<ztrowgraph_t>;
using crsAdapter_t = Zoltan2::TpetraCrsGraphAdapter<ztcrsgraph_t>;
using execspace_t =
    typename rowAdapter_t::ConstWeightsHostView1D::execution_space;
 
template <typename offset_t>
void printGraph(RCP<const Comm<int>> &comm, zlno_t nvtx, const zgno_t *vtxIds,
                const offset_t *offsets, const zgno_t *edgeIds) {
  int rank = comm->getRank();
  int nprocs = comm->getSize();
  comm->barrier();
  for (int p = 0; p < nprocs; p++) {
    if (p == rank) {
      std::cout << rank << ":" << std::endl;
      for (zlno_t i = 0; i < nvtx; i++) {
        std::cout << " vertex " << vtxIds[i] << ": ";
        for (offset_t j = offsets[i]; j < offsets[i + 1]; j++) {
          std::cout << edgeIds[j] << " ";
        }
        std::cout << std::endl;
      }
      std::cout.flush();
    }
    comm->barrier();
  }
  comm->barrier();
}
 
template <typename adapter_t, typename graph_t>
void TestGraphIds(adapter_t &ia, graph_t &graph) {
 
  using idsHost_t = typename adapter_t::ConstIdsHostView;
  using offsetsHost_t = typename adapter_t::ConstOffsetsHostView;
  using localInds_t =
      typename adapter_t::user_t::nonconst_local_inds_host_view_type;
 
  const auto nvtx = graph.getLocalNumRows();
  const auto nedges = graph.getLocalNumEntries();
  const auto maxNumEntries = graph.getLocalMaxNumRowEntries();
 
  typename adapter_t::Base::ConstIdsHostView adjIdsHost_("adjIdsHost_", nedges);
  typename adapter_t::Base::ConstOffsetsHostView offsHost_("offsHost_",
                                                           nvtx + 1);
 
  localInds_t nbors("nbors", maxNumEntries);
 
  for (size_t v = 0; v < nvtx; v++) {
    size_t numColInds = 0;
    graph.getLocalRowCopy(v, nbors, numColInds);
 
    offsHost_(v + 1) = offsHost_(v) + numColInds;
    for (size_t e = offsHost_(v), i = 0; e < offsHost_(v + 1); e++) {
      adjIdsHost_(e) = graph.getColMap()->getGlobalElement(nbors(i++));
    }
  }
 
  idsHost_t vtxIdsHost;
  ia.getVertexIDsHostView(vtxIdsHost);
 
  const auto graphIDS = graph.getRowMap()->getLocalElementList();
 
  Z2_TEST_COMPARE_ARRAYS(graphIDS, vtxIdsHost);
 
  idsHost_t adjIdsHost;
  offsetsHost_t offsetsHost;
  ia.getEdgesHostView(offsetsHost, adjIdsHost);
 
  Z2_TEST_COMPARE_ARRAYS(adjIdsHost_, adjIdsHost);
  Z2_TEST_COMPARE_ARRAYS(offsHost_, offsetsHost);
}
 
template <typename adapter_t, typename graph_t>
void verifyInputAdapter(adapter_t &ia, graph_t &graph) {
  using idsDevice_t = typename adapter_t::ConstIdsDeviceView;
  using idsHost_t = typename adapter_t::ConstIdsHostView;
  using offsetsDevice_t = typename adapter_t::ConstOffsetsDeviceView;
  using offsetsHost_t = typename adapter_t::ConstOffsetsHostView;
  using weightsDevice_t = typename adapter_t::WeightsDeviceView1D;
  using weightsHost_t = typename adapter_t::WeightsHostView1D;
 
  const auto nVtx = ia.getLocalNumIDs();
 
  Z2_TEST_EQUALITY(ia.getLocalNumVertices(), graph.getLocalNumRows());
  Z2_TEST_EQUALITY(ia.getLocalNumEdges(), graph.getLocalNumEntries());
 
 
  idsDevice_t vtxIdsDevice;
  ia.getVertexIDsDeviceView(vtxIdsDevice);
  idsHost_t vtxIdsHost;
  ia.getVertexIDsHostView(vtxIdsHost);
 
  Z2_TEST_DEVICE_HOST_VIEWS(vtxIdsDevice, vtxIdsHost);
 
 
  idsDevice_t adjIdsDevice;
  offsetsDevice_t offsetsDevice;
 
  ia.getEdgesDeviceView(offsetsDevice, adjIdsDevice);
 
  idsHost_t adjIdsHost;
  offsetsHost_t offsetsHost;
  ia.getEdgesHostView(offsetsHost, adjIdsHost);
 
  Z2_TEST_DEVICE_HOST_VIEWS(adjIdsDevice, adjIdsHost);
  Z2_TEST_DEVICE_HOST_VIEWS(offsetsDevice, offsetsHost);
 
  Z2_TEST_THROW(ia.setVertexWeightsDevice(
                    typename adapter_t::ConstWeightsDeviceView1D{}, 50),
                std::runtime_error);
 
  weightsDevice_t wgts0("wgts0", nVtx);
  Kokkos::parallel_for(
      nVtx, KOKKOS_LAMBDA(const int idx) { wgts0(idx) = idx * 2; });
  Kokkos::fence();
 
  Z2_TEST_NOTHROW(ia.setVertexWeightsDevice(wgts0, 0));
 
  // Don't reuse the same View, since we don't copy the values,
  // we just assign the View (increase use count)
  weightsDevice_t wgts1("wgts1", nVtx);
  Kokkos::parallel_for(
      nVtx, KOKKOS_LAMBDA(const int idx) { wgts1(idx) = idx * 3; });
 
  Z2_TEST_NOTHROW(ia.setVertexWeightsDevice(wgts1, 1));
 
  {
    weightsDevice_t weightsDevice;
    Z2_TEST_NOTHROW(ia.getVertexWeightsDeviceView(weightsDevice, 0));
 
    weightsHost_t weightsHost;
    Z2_TEST_NOTHROW(ia.getVertexWeightsHostView(weightsHost, 0));
 
    Z2_TEST_DEVICE_HOST_VIEWS(weightsDevice, weightsHost);
 
    Z2_TEST_DEVICE_HOST_VIEWS(wgts0, weightsHost);
  }
  {
    weightsDevice_t weightsDevice;
    Z2_TEST_NOTHROW(ia.getVertexWeightsDeviceView(weightsDevice, 1));
 
    weightsHost_t weightsHost;
    Z2_TEST_NOTHROW(ia.getVertexWeightsHostView(weightsHost, 1));
 
    Z2_TEST_DEVICE_HOST_VIEWS(weightsDevice, weightsHost);
 
    Z2_TEST_DEVICE_HOST_VIEWS(wgts1, weightsHost);
  }
  {
    weightsDevice_t wgtsDevice;
    Z2_TEST_THROW(ia.getVertexWeightsDeviceView(wgtsDevice, 2),
                  std::runtime_error);
 
    weightsHost_t wgtsHost;
    Z2_TEST_THROW(ia.getVertexWeightsHostView(wgtsHost, 2), std::runtime_error);
  }
 
  TestGraphIds(ia, graph);
}
 
int main(int narg, char *arg[]) {
  using rowSoln_t = Zoltan2::PartitioningSolution<rowAdapter_t>;
  using rowPart_t = rowAdapter_t::part_t;
 
  using crsSoln_t = Zoltan2::PartitioningSolution<crsAdapter_t>;
  using crsPart_t = crsAdapter_t::part_t;
 
  Tpetra::ScopeGuard tscope(&narg, &arg);
  const auto comm = Tpetra::getDefaultComm();
 
  try {
    Teuchos::ParameterList params;
    params.set("input file", "simple");
    params.set("file type", "Chaco");
 
    auto uinput = rcp(new UserInputForTests(params, comm));
 
    // Input crs graph and row graph cast from it.
    const auto crsGraph = uinput->getUITpetraCrsGraph();
    const auto rowGraph = rcp_dynamic_cast<ztrowgraph_t>(crsGraph);
 
    const auto nvtx = rowGraph->getLocalNumRows();
 
    // To test migration in the input adapter we need a Solution object.
    const auto env = rcp(new Zoltan2::Environment(comm));
 
    const int nWeights = 2;
 
    // User object is Tpetra::CrsGraph
    {
      PrintFromRoot("Input adapter for Tpetra::CrsGraph");
 
      auto tpetraCrsGraphInput = rcp(new crsAdapter_t(crsGraph, nWeights));
 
      verifyInputAdapter(*tpetraCrsGraphInput, *crsGraph);
 
      ztcrsgraph_t *mMigrate = NULL;
      crsPart_t *p = new crsPart_t[nvtx];
      memset(p, 0, sizeof(crsPart_t) * nvtx);
      ArrayRCP<crsPart_t> solnParts(p, 0, nvtx, true);
 
      crsSoln_t solution(env, comm, nWeights);
      solution.setParts(solnParts);
      tpetraCrsGraphInput->applyPartitioningSolution(*crsGraph, mMigrate,
                                                     solution);
      const auto newG = rcp(mMigrate);
 
      auto cnewG = rcp_const_cast<const ztcrsgraph_t>(newG);
      auto newInput = rcp(new crsAdapter_t(cnewG, nWeights));
 
      PrintFromRoot("Input adapter for Tpetra::RowGraph migrated to proc 0");
 
      verifyInputAdapter(*newInput, *newG);
    }
 
    // User object is Tpetra::RowGraph
    {
      PrintFromRoot("Input adapter for Tpetra::RowGraph");
 
      auto tpetraRowGraphInput = rcp(new rowAdapter_t(rowGraph, nWeights));
 
      verifyInputAdapter(*tpetraRowGraphInput, *crsGraph);
 
      rowPart_t *p = new rowPart_t[nvtx];
      memset(p, 0, sizeof(rowPart_t) * nvtx);
      ArrayRCP<rowPart_t> solnParts(p, 0, nvtx, true);
 
      rowSoln_t solution(env, comm, nWeights);
      solution.setParts(solnParts);
 
      ztrowgraph_t *mMigrate = NULL;
      tpetraRowGraphInput->applyPartitioningSolution(*crsGraph, mMigrate,
                                                     solution);
      const auto newG = rcp(mMigrate);
 
      auto cnewG = rcp_const_cast<const ztrowgraph_t>(newG);
      auto newInput = rcp(new rowAdapter_t(cnewG, nWeights));
 
      PrintFromRoot("Input adapter for Tpetra::RowGraph migrated to proc 0");
 
      verifyInputAdapter(*newInput, *newG);
    }
  } catch (std::exception &e) {
    std::cout << e.what() << std::endl;
    return EXIT_FAILURE;
  }
 
  PrintFromRoot("PASS");
}