Tpetra_BlockMultiVector_decl.hpp

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

#ifndef TPETRA_BLOCKMULTIVECTOR_DECL_HPP
#define TPETRA_BLOCKMULTIVECTOR_DECL_HPP

#include "Tpetra_BlockMultiVector_fwd.hpp"
#include "Tpetra_BlockCrsMatrix_fwd.hpp"
#include "Tpetra_MultiVector.hpp"
#include <memory>
#include <utility>

namespace Tpetra {

template <class Scalar,
          class LO,
          class GO,
          class Node>
class BlockMultiVector : public Tpetra::DistObject<Scalar, LO, GO, Node> {
 private:
  using dist_object_type = Tpetra::DistObject<Scalar, LO, GO, Node>;
  using STS              = Teuchos::ScalarTraits<Scalar>;
  using packet_type      = typename dist_object_type::packet_type;

 public:


  using map_type = Tpetra::Map<LO, GO, Node>;
  using mv_type = Tpetra::MultiVector<Scalar, LO, GO, Node>;

  using scalar_type = Scalar;
  using impl_scalar_type = typename mv_type::impl_scalar_type;
  using local_ordinal_type = LO;
  using global_ordinal_type = GO;
  using device_type = typename mv_type::device_type;

  using node_type = Node;

  using buffer_device_type = typename dist_object_type::buffer_device_type;

  typedef Kokkos::View<impl_scalar_type*, device_type> little_vec_type;
  typedef typename little_vec_type::host_mirror_type little_host_vec_type;

  typedef Kokkos::View<const impl_scalar_type*, device_type>
      const_little_vec_type;

  typedef typename mv_type::dual_view_type::t_host::device_type
      host_device_type;
  typedef Kokkos::View<const impl_scalar_type*, host_device_type>
      const_little_host_vec_type;



  BlockMultiVector();

  BlockMultiVector(const BlockMultiVector<Scalar, LO, GO, Node>&) = default;

  BlockMultiVector(BlockMultiVector<Scalar, LO, GO, Node>&&) = default;

  BlockMultiVector<Scalar, LO, GO, Node>&
  operator=(const BlockMultiVector<Scalar, LO, GO, Node>&) = default;

  BlockMultiVector<Scalar, LO, GO, Node>&
  operator=(BlockMultiVector<Scalar, LO, GO, Node>&&) = default;

  BlockMultiVector(const BlockMultiVector<Scalar, LO, GO, Node>& in,
                   const Teuchos::DataAccess copyOrView);

  BlockMultiVector(const map_type& meshMap,
                   const LO blockSize,
                   const LO numVecs);

  BlockMultiVector(const map_type& meshMap,
                   const map_type& pointMap,
                   const LO blockSize,
                   const LO numVecs);

  BlockMultiVector(const mv_type& X_mv,
                   const map_type& meshMap,
                   const LO blockSize);

  BlockMultiVector(const BlockMultiVector<Scalar, LO, GO, Node>& X,
                   const map_type& newMeshMap,
                   const map_type& newPointMap,
                   const size_t offset = 0);

  BlockMultiVector(const BlockMultiVector<Scalar, LO, GO, Node>& X,
                   const map_type& newMeshMap,
                   const size_t offset = 0);



  static map_type
  makePointMap(const map_type& meshMap, const LO blockSize);

  static Teuchos::RCP<const map_type>
  makePointMapRCP(const map_type& meshMap, const LO blockSize);

  const map_type getPointMap() const {
    return *pointMap_;
  }

  LO getBlockSize() const {
    return blockSize_;
  }

  LO getNumVectors() const {
    return static_cast<LO>(mv_.getNumVectors());
  }

  const mv_type& getMultiVectorView() const;
  mv_type& getMultiVectorView();



  void putScalar(const Scalar& val);

  void scale(const Scalar& val);

  void
  update(const Scalar& alpha,
         const BlockMultiVector<Scalar, LO, GO, Node>& X,
         const Scalar& beta);

  void
  blockWiseMultiply(const Scalar& alpha,
                    const Kokkos::View<const impl_scalar_type***,
                                       device_type, Kokkos::MemoryUnmanaged>& D,
                    const BlockMultiVector<Scalar, LO, GO, Node>& X);

  void
  blockJacobiUpdate(const Scalar& alpha,
                    const Kokkos::View<const impl_scalar_type***,
                                       device_type, Kokkos::MemoryUnmanaged>& D,
                    const BlockMultiVector<Scalar, LO, GO, Node>& X,
                    BlockMultiVector<Scalar, LO, GO, Node>& Z,
                    const Scalar& beta);

  template <class TargetMemorySpace>
  bool need_sync() const {
    return mv_.template need_sync<typename TargetMemorySpace::memory_space>();
  }

  bool need_sync_host() const {
    return mv_.need_sync_host();
  }

  bool need_sync_device() const {
    return mv_.need_sync_device();
  }



  bool replaceLocalValues(const LO localRowIndex, const LO colIndex, const Scalar vals[]);

  bool replaceGlobalValues(const GO globalRowIndex, const LO colIndex, const Scalar vals[]);

  bool sumIntoLocalValues(const LO localRowIndex, const LO colIndex, const Scalar vals[]);

  bool sumIntoGlobalValues(const GO globalRowIndex, const LO colIndex, const Scalar vals[]);

  const_little_host_vec_type getLocalBlockHost(
      const LO localRowIndex,
      const LO colIndex,
      const Access::ReadOnlyStruct) const;

  little_host_vec_type getLocalBlockHost(
      const LO localRowIndex,
      const LO colIndex,
      const Access::ReadWriteStruct);

  little_host_vec_type getLocalBlockHost(
      const LO localRowIndex,
      const LO colIndex,
      const Access::OverwriteAllStruct);

 protected:


  virtual bool checkSizes(const Tpetra::SrcDistObject& source) override;

  using dist_object_type::
      copyAndPermute;  

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

  using dist_object_type::packAndPrepare;  

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

  using dist_object_type::unpackAndCombine;  

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


 protected:
  size_t getStrideX() const {
    return static_cast<size_t>(1);
  }

  size_t getStrideY() const {
    return mv_.getStride();
  }

  bool isValidLocalMeshIndex(const LO meshLocalIndex) const;

  map_type meshMap_;

 private:
  Teuchos::RCP<const map_type> pointMap_;

 protected:
  mv_type mv_;

 private:
  LO blockSize_;

  void
  replaceLocalValuesImpl(const LO localRowIndex,
                         const LO colIndex,
                         const Scalar vals[]);
  void
  sumIntoLocalValuesImpl(const LO localRowIndex,
                         const LO colIndex,
                         const Scalar vals[]);

  static Teuchos::RCP<const mv_type>
  getMultiVectorFromSrcDistObject(const Tpetra::SrcDistObject&);

  static Teuchos::RCP<const BlockMultiVector<Scalar, LO, GO, Node> >
  getBlockMultiVectorFromSrcDistObject(const Tpetra::SrcDistObject& src);
};

}  // namespace Tpetra

#endif  // TPETRA_BLOCKMULTIVECTOR_DECL_HPP