Tpetra_Details_PackTriples.hpp

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef TPETRA_DETAILS_PACKTRIPLES_HPP
#define TPETRA_DETAILS_PACKTRIPLES_HPP

#include "TpetraCore_config.h"
#include "Teuchos_Comm.hpp"
#ifdef HAVE_TPETRACORE_MPI
#include "Tpetra_Details_extractMpiCommFromTeuchos.hpp"
#include "Tpetra_Details_MpiTypeTraits.hpp"
#endif  // HAVE_TPETRACORE_MPI
#include <ostream>

namespace Tpetra {
namespace Details {

//
// Search for "SKIP DOWN TO HERE" (omit quotes) for the "public"
// interface.  I put "public" in quotes because it's public only for
// Tpetra developers, NOT for Tpetra users.
//

#ifdef HAVE_TPETRACORE_MPI

// It's not useful to pack or unpack if not using MPI.  Thus, we only
// need to define these with MPI.  For a non-MPI build, we just need
// some outer stub interface that throws an exception.

int countPackTriplesCountMpi(MPI_Comm comm,
                             int& size,
                             std::ostream* errStrm = NULL);

int packTriplesCountMpi(const int numEnt,
                        char outBuf[],
                        const int outBufSize,
                        int& outBufCurPos,
                        MPI_Comm comm,
                        std::ostream* errStrm = NULL);

template <class ScalarType, class OrdinalType>
int countPackTriplesMpi(MPI_Datatype ordinalDt,
                        MPI_Datatype scalarDt,
                        const int numEnt,
                        MPI_Comm comm,
                        int& size,
                        std::ostream* errStrm = NULL) {
  using std::endl;
  int errCode = MPI_SUCCESS;

  int totalSize = 0;  // return value

  // Count the global row and column indices.
  {
    int curSize = 0;
    // We're packing two ordinals, the row resp. column index.
    errCode = MPI_Pack_size(2, ordinalDt, comm, &curSize);
    if (errCode != MPI_SUCCESS) {
      if (errStrm != NULL) {
        *errStrm << "countPackTripleMpi: MPI_Pack_size failed on "
                    "ordinalDt call"
                 << endl;
      }
      return errCode;
    }
    totalSize += curSize;
  }
  // Count the matrix value.
  {
    int curSize = 0;
    errCode     = MPI_Pack_size(1, scalarDt, comm, &curSize);
    if (errCode != MPI_SUCCESS) {
      if (errStrm != NULL) {
        *errStrm << "countPackTripleMpi: MPI_Pack_size failed on "
                    "scalarDt call"
                 << endl;
      }
      return errCode;
    }
    totalSize += curSize;  // one Scalar
  }
  totalSize *= numEnt;  // all the entries we want to pack

  size = totalSize;  // "commit" the result
  return errCode;
}

template <class ScalarType, class OrdinalType>
int packTripleMpi(const OrdinalType gblRowInd,
                  const OrdinalType gblColInd,
                  MPI_Datatype ordinalDt,
                  const ScalarType& val,
                  MPI_Datatype scalarDt,
                  char outBuf[],
                  const int outBufSize,
                  int& outBufCurPos,
                  MPI_Comm comm,
                  std::ostream* errStrm = NULL) {
  using std::endl;
  int errCode = MPI_SUCCESS;

  // Combine the two indices into a single MPI_Pack call.
  OrdinalType inds[2] = {gblRowInd, gblColInd};
  errCode             = MPI_Pack(inds, 2, ordinalDt,
                                 outBuf, outBufSize, &outBufCurPos, comm);
  if (errCode != MPI_SUCCESS) {
    if (errStrm != NULL) {
      *errStrm << "packTripleMpi: MPI_Pack failed for indices i=" << gblRowInd
               << ", j=" << gblColInd << ", where outBufSize=" << outBufSize
               << " and outBufCurPos=" << outBufCurPos << "." << endl;
    }
    return errCode;
  }
  // mfh 17,20 Jan 2017: Some (generally older) MPI implementations
  // want the first argument to be a pointer to nonconst, even though
  // MPI_Pack does not modify that argument.
  errCode = MPI_Pack(const_cast<ScalarType*>(&val), 1, scalarDt,
                     outBuf, outBufSize, &outBufCurPos, comm);
  if (errCode != MPI_SUCCESS) {
    if (errStrm != NULL) {
      *errStrm << "packTripleMpi: MPI_Pack failed on val = " << val
               << endl;
    }
    return errCode;
  }
  return errCode;
}

template <class ScalarType, class OrdinalType>
int packTriplesMpi(const OrdinalType gblRowInds[],
                   const OrdinalType gblColInds[],
                   MPI_Datatype ordinalDt,
                   const ScalarType vals[],
                   MPI_Datatype scalarDt,
                   const int numEnt,
                   char outBuf[],
                   const int outBufSize,
                   int& outBufCurPos,
                   MPI_Comm comm,
                   std::ostream* errStrm = NULL) {
  using std::endl;
  int errCode = MPI_SUCCESS;

  for (int k = 0; k < numEnt; ++k) {
    errCode = packTripleMpi(gblRowInds[k], gblColInds[k], ordinalDt,
                            vals[k], scalarDt, outBuf, outBufSize,
                            outBufCurPos, comm, errStrm);
    if (errCode != MPI_SUCCESS) {
      if (errStrm != NULL) {
        *errStrm << "packTriplesMpi: packTripleMpi failed at entry k=" << k
                 << endl;
      }
      return errCode;
    }
  }
  return errCode;
}

template <class ScalarType, class OrdinalType>
int unpackTripleMpi(const char inBuf[],
                    const int inBufSize,
                    int& inBufCurPos,
                    OrdinalType& gblRowInd,
                    OrdinalType& gblColInd,
                    MPI_Datatype ordinalDt,
                    ScalarType& val,
                    MPI_Datatype scalarDt,
                    MPI_Comm comm,
                    std::ostream* errStrm = NULL) {
  using std::endl;
  int errCode = MPI_SUCCESS;

  // mfh 17 Jan 2017: Some (generally older) MPI implementations want
  // the input buffer argument to be a pointer to nonconst, even
  // though MPI_Unpack does not modify that argument.

  // Combine the two indices into a single MPI_Pack call.
  OrdinalType inds[2] = {static_cast<OrdinalType>(0),
                         static_cast<OrdinalType>(0)};
  errCode             = MPI_Unpack(const_cast<char*>(inBuf), inBufSize, &inBufCurPos,
                                   inds, 2, ordinalDt, comm);
  if (errCode != MPI_SUCCESS) {
    if (errStrm != NULL) {
      *errStrm << "unpackTripleMpi: MPI_Unpack failed when unpacking indices: "
                  "inBufSize="
               << inBufSize << ", inBufCurPos=" << inBufCurPos << endl;
    }
    return errCode;
  }
  gblRowInd = inds[0];
  gblColInd = inds[1];

  // mfh 17 Jan 2017: Some (generally older) MPI implementations want
  // the input buffer argument to be a pointer to nonconst, even
  // though MPI_Unpack does not modify that argument.
  errCode = MPI_Unpack(const_cast<char*>(inBuf), inBufSize, &inBufCurPos,
                       &val, 1, scalarDt, comm);
  if (errCode != MPI_SUCCESS) {
    if (errStrm != NULL) {
      *errStrm << "unpackTripleMpi: MPI_Unpack failed when unpacking value: "
                  "inBufSize="
               << inBufSize << ", inBufCurPos=" << inBufCurPos << endl;
    }
    return errCode;
  }
  return errCode;
}

int unpackTriplesCountMpi(const char inBuf[],
                          const int inBufSize,
                          int& inBufCurPos,
                          int& numEnt,
                          MPI_Comm comm,
                          std::ostream* errStrm = NULL);

template <class ScalarType, class OrdinalType>
int unpackTriplesMpi(const char inBuf[],
                     const int inBufSize,
                     int& inBufCurPos,
                     OrdinalType gblRowInds[],
                     OrdinalType gblColInds[],
                     MPI_Datatype ordinalDt,
                     ScalarType vals[],
                     MPI_Datatype scalarDt,
                     const int numEnt,  // input arg, from unpackTriplesCountMpi
                     MPI_Comm comm,
                     std::ostream* errStrm = NULL) {
  using std::endl;
  int errCode = MPI_SUCCESS;

  for (int k = 0; k < numEnt; ++k) {
    errCode = unpackTripleMpi(inBuf, inBufSize, inBufCurPos,
                              gblRowInds[k], gblColInds[k], ordinalDt,
                              vals[k], scalarDt, comm, errStrm);
    if (errCode != MPI_SUCCESS) {
      if (errStrm != NULL) {
        *errStrm << "unpackTriplesMpi: packTripleMpi failed at entry k=" << k
                 << ": inBufSize=" << inBufSize
                 << ", inBufCurPos=" << inBufCurPos
                 << "." << endl;
      }
      return errCode;
    }
  }
  return errCode;
}

#endif  // HAVE_TPETRACORE_MPI

//
// SKIP DOWN TO HERE FOR "PUBLIC" INTERFACE
//

int countPackTriplesCount(const ::Teuchos::Comm<int>& comm,
                          int& size,
                          std::ostream* errStrm = NULL);

template <class ScalarType, class OrdinalType>
int countPackTriples(const int numEnt,
                     const ::Teuchos::Comm<int>& comm,
                     int& size,  // output argument
                     std::ostream* errStrm = NULL) {
#ifdef HAVE_TPETRACORE_MPI
  using ::Tpetra::Details::extractMpiCommFromTeuchos;
  using ::Tpetra::Details::MpiTypeTraits;

  static_assert(MpiTypeTraits<ScalarType>::isSpecialized,
                "countPackTriples: "
                "ScalarType lacks an MpiTypeTraits specialization.");
  static_assert(MpiTypeTraits<OrdinalType>::isSpecialized,
                "countPackTriples: "
                "OrdinalType lacks an MpiTypeTraits specialization.");

  MPI_Comm mpiComm       = extractMpiCommFromTeuchos(comm);
  MPI_Datatype ordinalDt = MpiTypeTraits<OrdinalType>::getType();
  MPI_Datatype scalarDt  = MpiTypeTraits<ScalarType>::getType();

  const int errCode =
      countPackTriplesMpi<ScalarType, OrdinalType>(ordinalDt, scalarDt,
                                                   numEnt, mpiComm,
                                                   size, errStrm);
  if (MpiTypeTraits<ScalarType>::needsFree) {
    (void)MPI_Type_free(&scalarDt);
  }
  if (MpiTypeTraits<OrdinalType>::needsFree) {
    (void)MPI_Type_free(&ordinalDt);
  }
  return errCode;

#else   // NOT HAVE_TPETRACORE_MPI
  if (errStrm != NULL) {
    *errStrm << "countPackTriples: Not implemented (no need; there's no need "
                "to pack or unpack anything if there's only one process)."
             << std::endl;
  }
  return -1;
#endif  // HAVE_TPETRACORE_MPI
}

int packTriplesCount(const int numEnt,
                     char outBuf[],
                     const int outBufSize,
                     int& outBufCurPos,
                     const ::Teuchos::Comm<int>& comm,
                     std::ostream* errStrm = NULL);

template <class ScalarType, class OrdinalType>
int
#ifdef HAVE_TPETRACORE_MPI
packTriples(const OrdinalType gblRowInds[],
            const OrdinalType gblColInds[],
            const ScalarType vals[],
            const int numEnt,
            char outBuf[],
            const int outBufSize,
            int& outBufCurPos,
            const ::Teuchos::Comm<int>& comm,
            std::ostream* errStrm = NULL)
#else   // NOT HAVE_TPETRACORE_MPI
packTriples(const OrdinalType /* gblRowInds */[],
            const OrdinalType /* gblColInds */[],
            const ScalarType /* vals */[],
            const int /* numEnt */,
            char /* outBuf */[],
            const int /* outBufSize */,
            int& /* outBufCurPos */,
            const ::Teuchos::Comm<int>& /* comm */,
            std::ostream* errStrm = NULL)
#endif  // HAVE_TPETRACORE_MPI
{
#ifdef HAVE_TPETRACORE_MPI
  using ::Tpetra::Details::extractMpiCommFromTeuchos;
  using ::Tpetra::Details::MpiTypeTraits;

  static_assert(MpiTypeTraits<ScalarType>::isSpecialized,
                "packTriples: "
                "ScalarType lacks an MpiTypeTraits specialization.");
  static_assert(MpiTypeTraits<OrdinalType>::isSpecialized,
                "packTriples: "
                "OrdinalType lacks an MpiTypeTraits specialization.");

  MPI_Comm mpiComm       = extractMpiCommFromTeuchos(comm);
  MPI_Datatype ordinalDt = MpiTypeTraits<OrdinalType>::getType();
  MPI_Datatype scalarDt  = MpiTypeTraits<ScalarType>::getType();

  const int errCode =
      packTriplesMpi(gblRowInds, gblColInds, ordinalDt, vals, scalarDt,
                     numEnt, outBuf, outBufSize, outBufCurPos, mpiComm, errStrm);
  if (MpiTypeTraits<ScalarType>::needsFree) {
    (void)MPI_Type_free(&scalarDt);
  }
  if (MpiTypeTraits<OrdinalType>::needsFree) {
    (void)MPI_Type_free(&ordinalDt);
  }
  return errCode;

#else   // NOT HAVE_TPETRACORE_MPI
  if (errStrm != NULL) {
    *errStrm << "packTriples: Not implemented (no need; there's no need to "
                "pack or unpack anything if there's only one process)."
             << std::endl;
  }
  return -1;
#endif  // HAVE_TPETRACORE_MPI
}

int unpackTriplesCount(const char inBuf[],
                       const int inBufSize,
                       int& inBufCurPos,
                       int& numEnt,  // output argument!
                       const ::Teuchos::Comm<int>& comm,
                       std::ostream* errStrm = NULL);

template <class ScalarType, class OrdinalType>
int
#ifdef HAVE_TPETRACORE_MPI
unpackTriples(const char inBuf[],
              const int inBufSize,
              int& inBufCurPos,
              OrdinalType gblRowInds[],
              OrdinalType gblColInds[],
              ScalarType vals[],
              const int numEnt,
              const ::Teuchos::Comm<int>& comm,
              std::ostream* errStrm = NULL)
#else   // NOT HAVE_TPETRACORE_MPI
unpackTriples(const char /* inBuf */[],
              const int /* inBufSize */,
              int& /* inBufCurPos */,
              OrdinalType /* gblRowInds */[],
              OrdinalType /* gblColInds */[],
              ScalarType /* vals */[],
              const int /* numEnt */,
              const ::Teuchos::Comm<int>& /* comm */,
              std::ostream* errStrm = NULL)
#endif  // HAVE_TPETRACORE_MPI
{
#ifdef HAVE_TPETRACORE_MPI
  using ::Tpetra::Details::extractMpiCommFromTeuchos;
  using ::Tpetra::Details::MpiTypeTraits;

  static_assert(MpiTypeTraits<ScalarType>::isSpecialized,
                "unpackTriples: "
                "ScalarType lacks an MpiTypeTraits specialization.");
  static_assert(MpiTypeTraits<OrdinalType>::isSpecialized,
                "unpackTriples: "
                "OrdinalType lacks an MpiTypeTraits specialization.");

  MPI_Comm mpiComm       = extractMpiCommFromTeuchos(comm);
  MPI_Datatype ordinalDt = MpiTypeTraits<OrdinalType>::getType();
  MPI_Datatype scalarDt  = MpiTypeTraits<ScalarType>::getType();

  const int errCode =
      unpackTriplesMpi(inBuf, inBufSize, inBufCurPos,
                       gblRowInds, gblColInds, ordinalDt,
                       vals, scalarDt, numEnt, mpiComm, errStrm);
  if (MpiTypeTraits<ScalarType>::needsFree) {
    (void)MPI_Type_free(&scalarDt);
  }
  if (MpiTypeTraits<OrdinalType>::needsFree) {
    (void)MPI_Type_free(&ordinalDt);
  }
  return errCode;

#else   // NOT HAVE_TPETRACORE_MPI
  if (errStrm != NULL) {
    *errStrm << "unpackTriples: Not implemented (no need; there's no need to "
                "pack or unpack anything if there's only one process)."
             << std::endl;
  }
  return -1;
#endif  // HAVE_TPETRACORE_MPI
}

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_PACKTRIPLES_HPP