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

// clang-format off
#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"
#include "Kokkos_ArithTraits.hpp"
#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:


    using packet_type = typename ::Kokkos::ArithTraits<Packet>::val_type;
    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);


  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);

  // clang-format on
  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);
  // clang-format off

    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