Tpetra_Distribution1D.hpp

// @HEADER
// ***********************************************************************
//
//          Tpetra: Templated Linear Algebra Services Package
//                 Copyright (2008) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Michael A. Heroux (maherou@sandia.gov)
//
// ************************************************************************
// @HEADER

//  1D Row-based Distribution class
//  Assumes square matrix
//  Karen Devine, SNL 
//

#ifndef __TPETRA_DISTRIBUTION1D_HPP
#define __TPETRA_DISTRIBUTION1D_HPP

namespace Tpetra {

// Forward definition
template <typename gno_t, typename scalar_t>
class DistributionLowerTriangularBlock;

template <typename gno_t, typename scalar_t>
class Distribution1D : public Distribution<gno_t,scalar_t> {
// 1D row-wise distribution of matrix and vector entries
// Rows and vector entries may be linearly or randomly distributed, or
// read from a file.
// Row map and vector map are identical

public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution<gno_t,scalar_t>::np;
  using Distribution<gno_t,scalar_t>::nrows;
  using Distribution<gno_t,scalar_t>::Mine;

  Distribution1D(size_t nrows_,
                 const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                 const Teuchos::ParameterList &params) :
                 Distribution<gno_t,scalar_t>(nrows_, comm_, params)
  {
    int npRow = -1;  // Number of processors among which to distribute rows;
                     // Will compute if not set by user
    const Teuchos::ParameterEntry *pe = params.getEntryPtr("nProcessorRows");
    if (pe != NULL) npRow = pe->getValue<int>(&npRow);

    TEUCHOS_TEST_FOR_EXCEPTION(npRow != -1 && npRow != np, std::logic_error,
                               " nProcessorRows " << npRow << " must equal" <<
                               " nProcessors " << np <<
                               " for 1D distribution");

    if (me == 0) std::cout << "\n 1D Distribution: "
                           << "\n     np     = " << np << std::endl;
  }

  // Return whether this rank owns vector entry i.
  virtual bool VecMine(gno_t i) = 0; 

  // Return whether this rank owns nonzero (i,j)
  // Vector map and row map are the same in 1D distribution.
  inline bool Mine(gno_t i, gno_t j) {return VecMine(i);}
  inline bool Mine(gno_t i, gno_t j, int p) {return VecMine(i);}
};

template <typename gno_t, typename scalar_t>
class Distribution1DLinear: public Distribution1D<gno_t,scalar_t> {

public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution<gno_t,scalar_t>::np;
  using Distribution<gno_t,scalar_t>::nrows;

  Distribution1DLinear(size_t nrows_,
                       const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                       const Teuchos::ParameterList &params) :
                       Distribution1D<gno_t,scalar_t>(nrows_, comm_, params)
  {
    gno_t nMyRows = getNumRow(me);
    myFirstRow = getFirstRow(me);
    myLastRow = myFirstRow + nMyRows - 1;
  }

  inline enum DistributionType DistType() { return OneDLinear; }

  inline bool VecMine(gno_t i) { return (i >= myFirstRow && i <= myLastRow); }

private:
  gno_t myFirstRow;
  gno_t myLastRow;

  inline size_t getNumRow(int p) { 
    return (nrows / np + (int(nrows % np) > p)); 
  }

  inline gno_t getFirstRow(int p) { 
    return (p * (nrows / np) + std::min<int>(int(nrows % np), p));
  }

// DistributionLowerTriangularBlock class needs a 1DLinear distribution
friend class DistributionLowerTriangularBlock<gno_t,scalar_t>;  
  
};

template <typename gno_t, typename scalar_t>
class Distribution1DRandom : public Distribution1D<gno_t,scalar_t> {


public:
  using Distribution<gno_t,scalar_t>::me;

  Distribution1DRandom(size_t nrows_,
                       const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                       const Teuchos::ParameterList &params) :
                       Distribution1D<gno_t,scalar_t>(nrows_, comm_, params) 
  { if (me == 0) std::cout << "    randomize = true" << std::endl; }

  inline enum DistributionType DistType() { return OneDRandom; }

  inline bool VecMine(gno_t i) { return (this->HashToProc(i) == me); }
};

template <typename gno_t, typename scalar_t>
class Distribution1DVec : public Distribution1D<gno_t,scalar_t> {
// Distribution of nonzeros is determined by the distribution of the
// vector entries, as read from a file.
//
// Assumptions include:
// -  Distribution file containing the vector part assignments (N lines) 
//    is provided.  This file is read during the constructor.  
//    Format for an NxN matrix:
//        line 1 to N:  0-based part assignment of vector entry

public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution<gno_t,scalar_t>::np;
  using Distribution<gno_t,scalar_t>::comm;
  using Distribution<gno_t,scalar_t>::nrows;
  using Distribution<gno_t,scalar_t>::Mine;

  Distribution1DVec(size_t nrows_,
                    const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                    const Teuchos::ParameterList &params,
                    std::string &distributionfile) :
                    Distribution1D<gno_t,scalar_t>(nrows_, comm_, params)
  {
    std::ifstream fpin;
    if (me == 0) {
      fpin.open(distributionfile.c_str(), std::ios::in);
      if (!fpin.is_open()) {
        std::cout << "Error:  distributionfile " << distributionfile 
             << " not found" << std::endl;
        exit(-1);
      }
    }

    // Read the vector part assignment and broadcast it to all processes.
    // Broadcast in chunks of bcastsize values.
    // TODO:  Make the vector part assignment more scalable instead of 
    // TODO:  storing entire vector on every process.

    vecpart = new int[nrows];

    const int bcastsize = 1000000;
    
    gno_t start = 0;
    int cnt = 0;
    for (size_t i = 0; i < nrows; i++) {
      if (me == 0) fpin >> vecpart[i];
      cnt++;
      if (cnt == bcastsize || i == nrows-1) {
        Teuchos::broadcast<int, int>(*comm, 0, cnt, &(vecpart[start]));
        start += cnt;
        cnt = 0;
      }
    }

    if (me == 0) fpin.close();
  }

  ~Distribution1DVec() {delete [] vecpart;}

  inline enum DistributionType DistType() { return OneDVec; }

  // Vector distribution was read in.
  inline bool VecMine(gno_t i) { return (vecpart[i] == me); }

protected:
  int *vecpart;             // part assignment of vector entries

};

}

#endif