Tpetra_DistObject_decl.hpp

Go to the documentation of this file.

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

#ifndef TPETRA_DISTOBJECT_DECL_HPP
#define TPETRA_DISTOBJECT_DECL_HPP


#include "Tpetra_Details_DistributorActor.hpp"
#include "Tpetra_Map.hpp"
#include "Tpetra_Import.hpp"
#include "Tpetra_Export.hpp"
#include "Tpetra_SrcDistObject.hpp"
#include "Tpetra_DistObject_fwd.hpp"
#if KOKKOS_VERSION >= 40799
#include "KokkosKernels_ArithTraits.hpp"
#else
#include "Kokkos_ArithTraits.hpp"
#endif
#include <memory>
#include <type_traits>

namespace KokkosClassic {
enum ReadWriteOption {
  ReadWrite    = 0, 
  OverwriteAll = 1  
};
}  // namespace KokkosClassic

namespace Tpetra {

template <class DistObjectType>
void removeEmptyProcessesInPlace(Teuchos::RCP<DistObjectType>& input,
                                 const Teuchos::RCP<const Map<typename DistObjectType::local_ordinal_type,
                                                              typename DistObjectType::global_ordinal_type,
                                                              typename DistObjectType::node_type> >& newMap);

template <class DistObjectType>
void removeEmptyProcessesInPlace(Teuchos::RCP<DistObjectType>& input);

template <class Packet,
          class LocalOrdinal,
          class GlobalOrdinal,
          class Node>
class DistObject : virtual public SrcDistObject,
                   virtual public Teuchos::Describable {
 public:


#if KOKKOS_VERSION >= 40799
  using packet_type = typename ::KokkosKernels::ArithTraits<Packet>::val_type;
#else
  using packet_type = typename ::Kokkos::ArithTraits<Packet>::val_type;
#endif
  using local_ordinal_type = LocalOrdinal;
  using global_ordinal_type = GlobalOrdinal;
  using node_type = Node;

  using device_type = typename Node::device_type;
  using execution_space = typename device_type::execution_space;

  using map_type = Map<local_ordinal_type, global_ordinal_type, node_type>;



  explicit DistObject(const Teuchos::RCP<const map_type>& map);

  DistObject(const DistObject<Packet, LocalOrdinal, GlobalOrdinal, Node>&) = default;

  DistObject& operator=(const DistObject<Packet, LocalOrdinal, GlobalOrdinal, Node>&) = default;

  DistObject(DistObject<Packet, LocalOrdinal, GlobalOrdinal, Node>&&) = default;

  DistObject& operator=(DistObject<Packet, LocalOrdinal, GlobalOrdinal, Node>&&) = default;

  virtual ~DistObject() = default;



  void
  doImport(const SrcDistObject& source,
           const Import<LocalOrdinal, GlobalOrdinal, Node>& importer,
           const CombineMode CM,
           const bool restrictedMode = false);

  void
  doExport(const SrcDistObject& source,
           const Export<LocalOrdinal, GlobalOrdinal, Node>& exporter,
           const CombineMode CM,
           const bool restrictedMode = false);

  void
  doImport(const SrcDistObject& source,
           const Export<LocalOrdinal, GlobalOrdinal, Node>& exporter,
           const CombineMode CM,
           const bool restrictedMode = false);

  void
  doExport(const SrcDistObject& source,
           const Import<LocalOrdinal, GlobalOrdinal, Node>& importer,
           const CombineMode CM,
           const bool restrictedMode = false);

  void
  beginImport(const SrcDistObject& source,
              const Import<LocalOrdinal, GlobalOrdinal, Node>& importer,
              const CombineMode CM,
              const bool restrictedMode = false);

  void
  beginExport(const SrcDistObject& source,
              const Export<LocalOrdinal, GlobalOrdinal, Node>& exporter,
              const CombineMode CM,
              const bool restrictedMode = false);

  void
  beginImport(const SrcDistObject& source,
              const Export<LocalOrdinal, GlobalOrdinal, Node>& exporter,
              const CombineMode CM,
              const bool restrictedMode = false);

  void
  beginExport(const SrcDistObject& source,
              const Import<LocalOrdinal, GlobalOrdinal, Node>& importer,
              const CombineMode CM,
              const bool restrictedMode = false);

  void
  endImport(const SrcDistObject& source,
            const Import<LocalOrdinal, GlobalOrdinal, Node>& importer,
            const CombineMode CM,
            const bool restrictedMode = false);

  void
  endExport(const SrcDistObject& source,
            const Export<LocalOrdinal, GlobalOrdinal, Node>& exporter,
            const CombineMode CM,
            const bool restrictedMode = false);

  void
  endImport(const SrcDistObject& source,
            const Export<LocalOrdinal, GlobalOrdinal, Node>& exporter,
            const CombineMode CM,
            const bool restrictedMode = false);

  void
  endExport(const SrcDistObject& source,
            const Import<LocalOrdinal, GlobalOrdinal, Node>& importer,
            const CombineMode CM,
            const bool restrictedMode = false);

  bool transferArrived() const;



  bool isDistributed() const;

  virtual Teuchos::RCP<const map_type> getMap() const { return map_; }



  void print(std::ostream& os) const;



  virtual std::string description() const;

  virtual void
  describe(Teuchos::FancyOStream& out,
           const Teuchos::EVerbosityLevel verbLevel =
               Teuchos::Describable::verbLevel_default) const;



  virtual void
  removeEmptyProcessesInPlace(const Teuchos::RCP<const map_type>& newMap);

  const Details::DistributorActor& getActor() { return distributorActor_; }


 protected:
  enum ReverseOption {
    DoForward,  //*!< Perform the transfer in forward mode.
    DoReverse   //*!< Perform the transfer in reverse mode.
  };

  virtual size_t constantNumberOfPackets() const;

  virtual void
  doTransfer(const SrcDistObject& src,
             const ::Tpetra::Details::Transfer<local_ordinal_type, global_ordinal_type, node_type>& transfer,
             const char modeString[],
             const ReverseOption revOp,
             const CombineMode CM,
             const bool restrictedMode);

  virtual bool
  reallocArraysForNumPacketsPerLid(const size_t numExportLIDs,
                                   const size_t numImportLIDs);

  using buffer_memory_space =
      ::Tpetra::Details::DefaultTypes::comm_buffer_memory_space<device_type>;

 public:
  using buffer_device_type =
      Kokkos::Device<typename device_type::execution_space,
                     buffer_memory_space>;

 protected:
  void beginTransfer(const SrcDistObject& src,
                     const ::Tpetra::Details::Transfer<local_ordinal_type, global_ordinal_type, node_type>& transfer,
                     const char modeString[],
                     const ReverseOption revOp,
                     const CombineMode CM,
                     const bool restrictedMode);

  void endTransfer(const SrcDistObject& src,
                   const ::Tpetra::Details::Transfer<local_ordinal_type, global_ordinal_type, node_type>& transfer,
                   const char modeString[],
                   const ReverseOption revOp,
                   const CombineMode CM,
                   const bool restrictedMode);

  void doPostRecvs(const Details::DistributorPlan& distributorPlan,
                   size_t constantNumPackets,
                   bool commOnHost,
                   std::shared_ptr<std::string> prefix,
                   const bool canTryAliasing,
                   const CombineMode CM);

  void doPostSends(const Details::DistributorPlan& distributorPlan,
                   size_t constantNumPackets,
                   bool commOnHost,
                   std::shared_ptr<std::string> prefix);

  void doPackAndPrepare(const SrcDistObject& src,
                        const Kokkos::DualView<const local_ordinal_type*, buffer_device_type>& exportLIDs,
                        size_t& constantNumPackets,
                        const execution_space& space);

  void doUnpackAndCombine(const Kokkos::DualView<const local_ordinal_type*, buffer_device_type>& remoteLIDs,
                          size_t constantNumPackets,
                          CombineMode CM,
                          const execution_space& space);



  virtual bool
  checkSizes(const SrcDistObject& source) = 0;

  virtual void
  copyAndPermute(const SrcDistObject& source,
                 const size_t numSameIDs,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& permuteToLIDs,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& permuteFromLIDs,
                 const CombineMode CM);

  virtual void copyAndPermute(
      const SrcDistObject& source, const size_t numSameIDs,
      const Kokkos::DualView<const local_ordinal_type*, buffer_device_type>& permuteToLIDs,
      const Kokkos::DualView<const local_ordinal_type*, buffer_device_type>& permuteFromLIDs,
      const CombineMode CM, const execution_space& space);

  virtual void
  packAndPrepare(const SrcDistObject& source,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& exportLIDs,
                 Kokkos::DualView<packet_type*,
                                  buffer_device_type>& exports,
                 Kokkos::DualView<size_t*,
                                  buffer_device_type>
                     numPacketsPerLID,
                 size_t& constantNumPackets);

  virtual void
  packAndPrepare(const SrcDistObject& source,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& exportLIDs,
                 Kokkos::DualView<packet_type*,
                                  buffer_device_type>& exports,
                 Kokkos::DualView<size_t*,
                                  buffer_device_type>
                     numPacketsPerLID,
                 size_t& constantNumPackets,
                 const execution_space& space);

  virtual void
  unpackAndCombine(const Kokkos::DualView<const local_ordinal_type*,
                                          buffer_device_type>& importLIDs,
                   Kokkos::DualView<packet_type*,
                                    buffer_device_type>
                       imports,
                   Kokkos::DualView<size_t*,
                                    buffer_device_type>
                       numPacketsPerLID,
                   const size_t constantNumPackets,
                   const CombineMode combineMode);

  virtual void
  unpackAndCombine(const Kokkos::DualView<const local_ordinal_type*,
                                          buffer_device_type>& importLIDs,
                   Kokkos::DualView<packet_type*,
                                    buffer_device_type>
                       imports,
                   Kokkos::DualView<size_t*,
                                    buffer_device_type>
                       numPacketsPerLID,
                   const size_t constantNumPackets,
                   const CombineMode combineMode,
                   const execution_space& space);

  Teuchos::RCP<const map_type> map_;

 protected:
  std::unique_ptr<std::string>
  createPrefix(const char className[],
               const char methodName[]) const;

  Kokkos::DualView<packet_type*, buffer_device_type> imports_;

  virtual bool
  reallocImportsIfNeeded(const size_t newSize,
                         const bool verbose,
                         const std::string* prefix,
                         const bool remoteLIDsContiguous = false,
                         const CombineMode CM            = INSERT);

  Kokkos::DualView<size_t*, buffer_device_type> numImportPacketsPerLID_;

  Kokkos::DualView<packet_type*, buffer_device_type> exports_;

  Kokkos::DualView<size_t*, buffer_device_type> numExportPacketsPerLID_;

 private:
  using this_type = DistObject<Packet, LocalOrdinal, GlobalOrdinal, Node>;

  Details::DistributorActor distributorActor_;

};  // class DistObject
}  // namespace Tpetra

#endif  // TPETRA_DISTOBJECT_DECL_HPP