Zoltan2_MatrixPartitioningAlgs.hpp

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

#ifndef _ZOLTAN2_ALGMATRIX_HPP_
#define _ZOLTAN2_ALGMATRIX_HPP_

// #include <Zoltan2_IdentifierModel.hpp>
#include <Zoltan2_PartitioningSolution.hpp>
#include <Zoltan2_MatrixPartitioningSolution.hpp>
// #include <Zoltan2_Algorithm.hpp>

#include <Zoltan2_PartitioningProblem.hpp>

// #include <sstream>
// #include <string>
// #include <bitset>

namespace Zoltan2{


template <typename Adapter>
class AlgMatrix : public Algorithm<Adapter>
{

public:
  typedef typename Adapter::lno_t lno_t;     // local ids
  typedef typename Adapter::gno_t gno_t;     // global ids
  typedef typename Adapter::scalar_t scalar_t;   // scalars
  typedef typename Adapter::part_t part_t;   // part numbers

  typedef typename Adapter::user_t user_t;
  typedef typename Adapter::userCoord_t userCoord_t;

  typedef typename Adapter::base_adapter_t base_adapter_t;



  // Constructor
  AlgMatrix(const RCP<const Environment> &_env,
        const RCP<const Comm<int> > &_problemComm,
        const RCP<const MatrixAdapter<user_t, userCoord_t> > &_matrixAdapter)
    : mEnv(_env), mProblemComm(_problemComm), mMatrixAdapter(_matrixAdapter)
  {}


  // AlgMatrix(const RCP<const Environment> &_env,
  //      const RCP<const Comm<int> > &_problemComm,
  //      const RCP<const Adapter> &_matrixAdapter)
  //   : mEnv(_env), mProblemComm(_problemComm), mMatrixAdapter(_matrixAdapter)
  // {}

  // Partitioning method
  void partitionMatrix(const RCP<MatrixPartitioningSolution<Adapter> > &solution);


private:
  const RCP<const Environment> mEnv;
  const RCP<const Comm<int> > mProblemComm;

  const RCP<const MatrixAdapter<user_t, userCoord_t> > mMatrixAdapter;
  //const RCP<const Adapter > mMatrixAdapter;

  //typedef Tpetra::CrsMatrix<z2TestScalar, z2TestLO, z2TestGO> SparseMatrix;
  //  const RCP<const XpetraCrsMatrixAdapter<user_t, userCoord_t> > mMatrixAdapter;




};


// Partitioning method
//
//   -- For now implementing 2D Random Cartesian
template <typename Adapter>
void AlgMatrix<Adapter>::partitionMatrix(const RCP<MatrixPartitioningSolution<Adapter> > &solution)
{
  mEnv->debug(DETAILED_STATUS, std::string("Entering AlgBlock"));

  int myrank = mEnv->myRank_;
  int nprocs = mEnv->numProcs_;

  // Determine processor dimension d for 2D block partitioning d = sqrt(p) 
  // Determine processor row and processor column for this rank for 2D partitioning
  int procDim = sqrt(nprocs);
  assert(procDim * procDim == nprocs); // TODO: Should check this earlier and produce more useful error
  
  int myProcRow = myrank / procDim;
  int myProcCol = myrank % procDim;

  // Create 1D Random partitioning problem to create partitioning for the 2D partitioning

  // Create parameters for 1D partitioning
  Teuchos::ParameterList params1D("Params for 1D partitioning");
  //  params1D.set("algorithm", "random"); //TODO add support for random
  params1D.set("algorithm", "block");
  params1D.set("imbalance_tolerance", 1.1);
  params1D.set("num_global_parts", procDim);

  // Create Zoltan2 partitioning problem 
  //    -- Alternatively we could simply call algorithm directly
  MatrixAdapter<user_t, userCoord_t> *adapterPtr = 
    const_cast<MatrixAdapter<user_t, userCoord_t>*>(mMatrixAdapter.getRawPtr());
  

  Zoltan2::PartitioningProblem<base_adapter_t> *problem = 
    new Zoltan2::PartitioningProblem<base_adapter_t>(adapterPtr, &params1D);


  // Solve the problem, get the solution
  problem->solve();

  //  Zoltan2::PartitioningSolution<SparseMatrixAdapter_t> solution = problem.getSolution();
  const Zoltan2::PartitioningSolution<base_adapter_t> &solution1D = problem->getSolution();


  // Get part assignments for each local_id
  //  size_t numGlobalParts = solution1D->getTargetGlobalNumberOfParts();
  
  const part_t *parts = solution1D.getPartListView();




  // Create column Ids ArrayRCP colIDs to store in solution
  //    This will define which column IDs are allowed in a particular processor
  //    column.

  // Group gids corresponding to local ids based on process column 
  const gno_t *rowIds;
  mMatrixAdapter->getRowIDsView(rowIds);

  size_t nLocIDs = mMatrixAdapter->getLocalNumRows();

  std::vector<std::vector<gno_t> > idsInProcColI(procDim);
  Teuchos::ArrayRCP<gno_t> colIDs;

  for(size_t i=0; i<nLocIDs; i++)
  {
    // Nonzeros with columns of index rowIds[i] belong to some processor
    // in processor column parts[i]
    idsInProcColI[ parts[i] ].push_back(rowIds[i]);
  }

  delete problem; // delete 1D partitioning problem


  // Communicate gids to process col roots 

  // Loop over each of the processor columns
  for(int procColI=0; procColI<procDim; procColI++)
  {

    // This rank is root of procColI, need to receive
    if(myProcCol==procColI && myProcRow==0) 
    {
      assert(myrank==procColI); // Could make this above conditional?


      std::set<gno_t> gidSet; // to get sorted list of gids

      // Copy local gids to sorted gid set
      for(int i=0;i<idsInProcColI[procColI].size();i++)
      {
        gidSet.insert(idsInProcColI[procColI][i]);
      }

      // Receive gids from remote processors, insert into set
      std::vector<gno_t> recvBuf;
      for(int src=0;src<nprocs;src++)
      {
        if(src!=myrank)
    {
          int buffSize;

      Teuchos::receive<int,int>(*mProblemComm, src, 1, &buffSize);

          if(buffSize>0)
      {
            recvBuf.resize(buffSize);

        Teuchos::receive<int,gno_t>(*mProblemComm, src, buffSize, recvBuf.data());

            for(int i=0;i<buffSize;i++)
        {
              gidSet.insert(recvBuf[i]);
        }
      }
    }
      }

      //Copy data to std::vector 
      colIDs.resize(gidSet.size());

      typename std::set<gno_t>::const_iterator iter;
      int indx=0;
      for (iter=gidSet.begin(); iter!=gidSet.end(); ++iter)
      {
        colIDs[indx] = *iter;
    indx++;
      }
    }

    // This rank is not root of procColI, need to senddata block to root
    else
    {
      int sizeToSend = idsInProcColI[procColI].size();

      //is the dst proc info correct here?

      // Root of procColI is rank procColI
      Teuchos::send<int,int>(*mProblemComm,1,&sizeToSend,procColI);

      Teuchos::send<int,gno_t>(*mProblemComm,sizeToSend,idsInProcColI[procColI].data(),procColI);
    }

  }

  // Free memory if possible


  // Communicate gids from process col root down process column to other procs 

  // This rank is root of its processor column, need to send
  if(myProcRow==0) 
  {
    // Send data to remote processors in processor column
    for(int procColRank=0; procColRank<procDim; procColRank++)
    {
      // Convert from proc column rank to global rank
      int dstRank = procColRank*procDim + myProcCol;

      if(dstRank!=myrank)
      {
        int sizeToSend = colIDs.size();

    Teuchos::send<int,int>(*mProblemComm,1,&sizeToSend,dstRank);
        Teuchos::send<int,gno_t>(*mProblemComm,sizeToSend,colIDs.get(),dstRank);

      }
    }

  }

  // This rank is not root of processor, need to recv data from root
  else
  {
    // Src rank is root of processor column = myProcCol
    int srcRank = myProcCol;

    int buffSize;
    Teuchos::receive<int,int>(*mProblemComm, srcRank, 1, &buffSize);

    colIDs.resize(buffSize);
    Teuchos::receive<int,gno_t>(*mProblemComm, srcRank, buffSize, colIDs.get());
  }

  // Create domain/range IDs (same for both now) Array RCP to store in solution
  //   Created from equal division of column IDs

  // Split processor column ids into nDim parts
  // Processor column i ids split between ranks 0+i, nDim+i, 2*nDim+i, ...

  // 

  ArrayRCP<gno_t> domRangeIDs;     // Domain/Range vector Ids assigned to this process

  size_t nCols = colIDs.size();
  lno_t nColsDivDim = nCols / procDim;
  lno_t nColsModDim = nCols % procDim;

  size_t sColIndx;
  size_t domRangeSize;

  // This proc will have nColsDivDim+1 domain/range ids
  if(myProcRow < nColsModDim)
  {
    sColIndx = myProcRow * (nColsDivDim+1);
    domRangeSize = nColsDivDim+1;
  }
  // This proc will have nColsDivDim domain/range ids
  else
  {
    sColIndx =  nColsModDim*(nColsDivDim+1) + (myProcRow-nColsModDim) * nColsDivDim;
    domRangeSize = nColsDivDim;
  }

  domRangeIDs.resize(domRangeSize);

  for(size_t i=0;i<domRangeSize;i++)
  {
    domRangeIDs[i] = colIDs[sColIndx+i];
  }  

  // Create row IDs Array RCP to store in solution
  //   Created from union of domRangeIDs
  //
  // Procs 0, 1, ... nDim-1 share the same set of rowIDs
  // Procs nDim, nDim+1, ..., 2*nDim-1 share the same set of rowIDs

  // Create subcommunicator for processor row
  std::vector<int> subcommRanks(procDim);

  for(unsigned int i=0; i<procDim; i++)
  {
    subcommRanks[i] = myProcRow*procDim + i;
  }

  ArrayView<const int> subcommRanksView = Teuchos::arrayViewFromVector (subcommRanks);

  RCP<Teuchos::Comm<int> > rowcomm = mProblemComm->createSubcommunicator(subcommRanksView);

  // Determine max number of columns in this process row
  size_t maxProcRowNCols=0;

  Teuchos::reduceAll<int, size_t>(*rowcomm,Teuchos::REDUCE_MAX,1,&domRangeSize,&maxProcRowNCols);

  // Communicate all domRangeIDs to all processes in procRow
  gno_t MAXVAL = std::numeric_limits<gno_t>::max();

  std::vector<gno_t> locRowIDs(maxProcRowNCols,MAXVAL);

  Teuchos::ArrayRCP<gno_t> rowIDs(maxProcRowNCols * procDim);

  // Copy local domRangeIDs into local row IDs, "extra elements" will have
  // value MAXVAL
  for(size_t i=0;i<domRangeIDs.size();i++)
  {
    locRowIDs[i] = domRangeIDs[i];
  }    

  // Gather all ids onto all processes in procRow
  Teuchos::gatherAll<int,gno_t>(*rowcomm,maxProcRowNCols, locRowIDs.data(),
                                maxProcRowNCols*procDim, rowIDs.get());

  // Free local data
  std::vector<int>().swap(locRowIDs);


  // Insert local domRangeIDs into row IDs set, filter out extra items
  std::set<gno_t> setRowIDs;

  for(size_t i=0;i<rowIDs.size();i++)
  {
    if(rowIDs[i]!=MAXVAL)
    {
      setRowIDs.insert(rowIDs[i]);
    }
  }  

  // Resize rowIDs array and copy data to array (now sorted)
  rowIDs.resize(setRowIDs.size());

  typename std::set<gno_t>::const_iterator iter;
  size_t indx=0;

  for(iter=setRowIDs.begin(); iter!=setRowIDs.end(); ++iter)
  {
    rowIDs[indx] = *iter;
    indx++;
  }  


  // Finished. Store partition in solution.
  solution->setIDLists(rowIDs,colIDs,domRangeIDs,domRangeIDs);

  mEnv->debug(DETAILED_STATUS, std::string("Exiting AlgMatrix"));
}



}   // namespace Zoltan2

#endif