TaskMappingTest.cpp

Go to the documentation of this file.

#include "Zoltan2_TaskMapping.hpp"
#include <Zoltan2_PartitioningProblem.hpp>
#include <Zoltan2_TestHelpers.hpp>

#include <string>

#include <Teuchos_RCP.hpp>
#include <Teuchos_Array.hpp>
#include <Teuchos_ParameterList.hpp>
#include "Teuchos_XMLParameterListHelpers.hpp"

#include "Tpetra_MultiVector.hpp"
#include <Tpetra_CrsGraph.hpp>
#include <Tpetra_Map.hpp>

#include <Zoltan2_XpetraCrsGraphAdapter.hpp>
#include <Zoltan2_XpetraMultiVectorAdapter.hpp>

typedef Tpetra::CrsGraph<zlno_t, zgno_t, znode_t> tcrsGraph_t;
typedef Tpetra::MultiVector<zscalar_t, zlno_t, zgno_t, znode_t> tMVector_t;
typedef Zoltan2::XpetraCrsGraphAdapter<tcrsGraph_t, tMVector_t> my_adapter_t;


int main(int narg, char *arg[]){

  Tpetra::ScopeGuard tscope(&narg, &arg);
  Teuchos::RCP<const Teuchos::Comm<int> > tcomm = Tpetra::getDefaultComm();

  typedef my_adapter_t::part_t part_t;

  int nx = 2, ny = 2, nz = 2;
  for ( int i = 1 ; i < narg ; ++i ) {
    if ( 0 == strcasecmp( arg[i] , "NX" ) ) {
      nx = atoi( arg[++i] );
    }
    else if ( 0 == strcasecmp( arg[i] , "NY" ) ) {
      ny = atoi( arg[++i] );
    }
    else if ( 0 == strcasecmp( arg[i] , "NZ" ) ) {
      nz = atoi( arg[++i] );
    }
    else{
      std::cerr << "Unrecognized command line argument #" << i << ": " << arg[i] << std::endl ;
      return 1;
    }
  }


  try{


    int rank = tcomm->getRank();
    part_t numProcs = tcomm->getSize();

    int coordDim = 3;
    zgno_t numGlobalTasks = nx*ny*nz;

    zgno_t myTasks = numGlobalTasks / numProcs;
    zgno_t taskLeftOver = numGlobalTasks % numProcs;
    if (rank < taskLeftOver ) ++myTasks;

    zgno_t myTaskBegin = (numGlobalTasks / numProcs) * rank;
    myTaskBegin += (taskLeftOver < rank ? taskLeftOver : rank);
    zgno_t myTaskEnd = myTaskBegin + myTasks;

    zscalar_t **partCenters = NULL;
    partCenters = new zscalar_t * [coordDim];
    for(int i = 0; i < coordDim; ++i){
      partCenters[i] = new zscalar_t[myTasks];
    }


    zgno_t *task_gnos = new zgno_t [myTasks];
    zlno_t *task_communication_xadj_ = new zlno_t [myTasks+1];
    zgno_t *task_communication_adj_ = new zgno_t [myTasks * 6];

    zlno_t prevNCount = 0;
    task_communication_xadj_[0] = 0;
    for (zlno_t i = myTaskBegin; i < myTaskEnd; ++i) {
      task_gnos[i - myTaskBegin] = i;

      zlno_t x = i % nx;
      zlno_t y = (i / (nx)) % ny;
      zlno_t z = (i / (nx)) / ny;
      partCenters[0][i - myTaskBegin] = x;
      partCenters[1][i - myTaskBegin] = y;
      partCenters[2][i - myTaskBegin] = z;



      if (x > 0){
        task_communication_adj_[prevNCount++] = i - 1;
      }
      if (x < nx - 1){
        task_communication_adj_[prevNCount++] = i + 1;
      }
      if (y > 0){
        task_communication_adj_[prevNCount++] = i - nx;
      }
      if (y < ny - 1){
        task_communication_adj_[prevNCount++] = i + nx;
      }
      if (z > 0){
        task_communication_adj_[prevNCount++] = i - nx * ny;
      }
      if (z < nz - 1){
        task_communication_adj_[prevNCount++] = i + nx * ny;
      }
      task_communication_xadj_[i+1 - myTaskBegin] = prevNCount;
    }
    using namespace Teuchos;
    RCP<my_adapter_t> ia;
    typedef Tpetra::Map<>::node_type mytest_znode_t;
    typedef Tpetra::Map<zlno_t, zgno_t, mytest_znode_t> map_t;
    RCP<const map_t> map = rcp (new map_t (numGlobalTasks, myTasks, 0, tcomm));

    RCP<tcrsGraph_t> TpetraCrsGraph(new tcrsGraph_t (map, 0));


    for (zlno_t lclRow = 0; lclRow < myTasks; ++lclRow) {
      const zgno_t gblRow = map->getGlobalElement (lclRow);
      zgno_t begin = task_communication_xadj_[lclRow];
      zgno_t end = task_communication_xadj_[lclRow + 1];
      const ArrayView< const zgno_t > indices(task_communication_adj_+begin, end-begin);
      TpetraCrsGraph->insertGlobalIndices(gblRow, indices);
    }
    TpetraCrsGraph->fillComplete ();
    RCP<const tcrsGraph_t> const_data = rcp_const_cast<const tcrsGraph_t>(TpetraCrsGraph);

    ia = RCP<my_adapter_t> (new my_adapter_t(const_data));

    const int coord_dim = 3;
    Teuchos::Array<Teuchos::ArrayView<const zscalar_t> > coordView(coord_dim);

    if(myTasks > 0){
      Teuchos::ArrayView<const zscalar_t> a(partCenters[0], myTasks);
      coordView[0] = a;
      Teuchos::ArrayView<const zscalar_t> b(partCenters[1], myTasks);
      coordView[1] = b;
      Teuchos::ArrayView<const zscalar_t> c(partCenters[2], myTasks);
      coordView[2] = c;
    }
    else {
      Teuchos::ArrayView<const zscalar_t> a;
      coordView[0] = a;
      coordView[1] = a;
      coordView[2] = a;
    }
    RCP<tMVector_t> coords(new tMVector_t(map, coordView.view(0, coord_dim), coord_dim));//= set multivector;
    RCP<const tMVector_t> const_coords = rcp_const_cast<const tMVector_t>(coords);
    RCP <Zoltan2::XpetraMultiVectorAdapter<tMVector_t> > adapter (new Zoltan2::XpetraMultiVectorAdapter<tMVector_t>(const_coords));
    ia->setCoordinateInput(adapter.getRawPtr());


    //return ia;




    //create input adapter
    //RCP<my_adapter_t> ia = create_problem(tcomm, nx, ny, nz);

    //create partitioning problem
    typedef Zoltan2::PartitioningProblem<my_adapter_t> xcrsGraph_problem_t; // xpetra_graph problem type
    typedef Zoltan2::EvaluatePartition<my_adapter_t> quality_t;
    ParameterList zoltan2_parameters;
    zoltan2_parameters.set("compute_metrics", true); // bool parameter
    zoltan2_parameters.set("imbalance_tolerance", 1.0);
    zoltan2_parameters.set("num_global_parts", tcomm->getSize());
    zoltan2_parameters.set("algorithm", "multijagged");
    zoltan2_parameters.set("mj_keep_part_boxes", false); // bool parameter
    zoltan2_parameters.set("mj_recursion_depth", 3);
    RCP<xcrsGraph_problem_t> partition_problem;
    partition_problem  = RCP<xcrsGraph_problem_t> (new xcrsGraph_problem_t(ia.getRawPtr(),&zoltan2_parameters,tcomm));

    //solve the partitioning problem.
    partition_problem->solve();
    tcomm->barrier();
    RCP<const Zoltan2::Environment> env = partition_problem->getEnvironment();

    RCP<quality_t>metricObject = 
      rcp(new quality_t(ia.getRawPtr(), &zoltan2_parameters, tcomm,
      &partition_problem->getSolution()));

    if (tcomm->getRank() == 0){
      metricObject->printMetrics(std::cout);
    }
    partition_problem->printTimers();

    part_t *proc_to_task_xadj_ = NULL;
    part_t *proc_to_task_adj_ = NULL;

    //create the zoltan2 machine representation object
    Zoltan2::MachineRepresentation<zscalar_t, part_t> mach(*tcomm);

    //create the mapper and map the partitioning solution.
    Zoltan2::CoordinateTaskMapper<my_adapter_t, part_t> ctm (
        tcomm,
        Teuchos::rcpFromRef(mach),
        ia,
        rcpFromRef(partition_problem->getSolution()),
        env);

    //get the results and print
    ctm.getProcTask(proc_to_task_xadj_, proc_to_task_adj_);
    //part_t numProcs = tcomm->getSize();
    if (tcomm->getRank() == 0){
      for (part_t i = 0; i < numProcs; ++i){
        std::cout << "\nProc i:" << i << " ";
        for (part_t j = proc_to_task_xadj_[i]; j < proc_to_task_xadj_[i+1]; ++j){
          std::cout << " " << proc_to_task_adj_[j];
        }
      }
      std::cout << std::endl;
    }

    //below is to calculate the result hops. this uses the global graph
    //also this is used for debug, as the hops are also calculated in mapper.
    {

      zlno_t prevNCount_tmp = 0;
      zgno_t *task_communication_adj_tmp = new zgno_t [numGlobalTasks * 6];
      zlno_t *task_communication_xadj_tmp = new zlno_t [numGlobalTasks+1];
      task_communication_xadj_tmp[0] = 0;
      for (zlno_t i = 0; i < numGlobalTasks; ++i) {
        zlno_t x = i % nx;
        zlno_t y = (i / (nx)) % ny;
        zlno_t z = (i / (nx)) / ny;

        if (x > 0){
          task_communication_adj_tmp[prevNCount_tmp++] = i - 1;
        }
        if (x < nx - 1){
          task_communication_adj_tmp[prevNCount_tmp++] = i + 1;
        }
        if (y > 0){
          task_communication_adj_tmp[prevNCount_tmp++] = i - nx;
        }
        if (y < ny - 1){
          task_communication_adj_tmp[prevNCount_tmp++] = i + nx;
        }
        if (z > 0){
          task_communication_adj_tmp[prevNCount_tmp++] = i - nx * ny;
        }
        if (z < nz - 1){
          task_communication_adj_tmp[prevNCount_tmp++] = i + nx * ny;
        }
        task_communication_xadj_tmp[i+1] = prevNCount_tmp;
      }

      int mach_coord_dim = mach.getMachineDim();
      std::vector <int> mach_extent(mach_coord_dim);
      mach.getMachineExtent(&(mach_extent[0]));

      std::vector <part_t> all_parts(numGlobalTasks), copy(numGlobalTasks, 0);

      const part_t *parts = partition_problem->getSolution().getPartListView();

      //typedef Tpetra::Map<>::node_type mytest_znode_t;
      //typedef Tpetra::Map<zlno_t, zgno_t, mytest_znode_t> map_t;
      //RCP<const map_t> map = rcp (new map_t (numGlobalTasks, myTasks, 0, tcomm));
      for (part_t i = 0; i < myTasks; ++i){
        zgno_t g = map->getGlobalElement(i);
        copy[g] = parts[i];
      }

      reduceAll<int, part_t>(
          *tcomm,
          Teuchos::REDUCE_SUM,
          numGlobalTasks,
          &(copy[0]),
          &(all_parts[0])
      );

      zscalar_t **proc_coords;
      mach.getAllMachineCoordinatesView(proc_coords);
      part_t hops=0;
      part_t hops2 = 0;
      int *machine_extent = new int [mach_coord_dim];
      bool *machine_extent_wrap_around = new bool[mach_coord_dim];

      mach.getMachineExtent(machine_extent);
      mach.getMachineExtentWrapArounds(machine_extent_wrap_around);

      for (zlno_t i = 0; i < numGlobalTasks; ++i){
        zlno_t b = task_communication_xadj_tmp[i];
        zlno_t e = task_communication_xadj_tmp[i+1];

        part_t procId1 = ctm.getAssignedProcForTask(all_parts[i]);

        for (zlno_t j = b; j < e; ++j){
          zgno_t n = task_communication_adj_tmp[j];
          part_t procId2 = ctm.getAssignedProcForTask(all_parts[n]);

          zscalar_t distance2 = 0;
          mach.getHopCount(procId1, procId2, distance2);
          hops2 += distance2;
          for (int k = 0 ; k < mach_coord_dim ; ++k){
            part_t distance = ZOLTAN2_ABS(proc_coords[k][procId1] - proc_coords[k][procId2]);
            if (machine_extent_wrap_around[k]){
              if (machine_extent[k] - distance < distance){
                distance = machine_extent[k] - distance;
              }
            }
            hops += distance;
          }
        }
      }
      delete [] machine_extent_wrap_around;
      delete [] machine_extent;

      if (tcomm->getRank() == 0)
        std::cout << "HOPS:" << hops << " HOPS2:" << hops2 << std::endl;

      delete [] task_communication_xadj_tmp;
      delete [] task_communication_adj_tmp;
    }
    /*
    part_t *machineDimensions = NULL;
    //machineDimensions = hopper;
    Zoltan2::coordinateTaskMapperInterface<part_t, zscalar_t, zscalar_t>(
        tcomm,
        procDim,
        numProcs,
        procCoordinates,

        coordDim,
        numGlobalTasks,
        partCenters,

        task_communication_xadj_,
        task_communication_adj_,
        NULL,

        proc_to_task_xadj_,
        proc_to_task_adj_,

        partArraysize,
        partArray,
        machineDimensions
    );
     */

    if (tcomm->getRank() == 0){
      std::cout << "PASS" << std::endl;
    }


    for (int i = 0; i < coordDim; i++) delete [] partCenters[i];
    delete [] partCenters;

  }
  catch(std::string &s){
    std::cerr << s << std::endl;
  }

  catch(char * s){
    std::cerr << s << std::endl;
  }
}