Tpetra_idot.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_DETAILS_IDOT_HPP
#define TPETRA_DETAILS_IDOT_HPP


#include "Tpetra_iallreduce.hpp"
#include "Tpetra_MultiVector.hpp"
#include "Tpetra_Vector.hpp"
#include "Teuchos_CommHelpers.hpp"
#include "KokkosBlas1_dot.hpp"
#include <stdexcept>
#include <sstream>

namespace Tpetra {
namespace Details {

// Temporary helper to get read-only view from multivector in requested space.
// TODO: when https://github.com/kokkos/kokkos/issues/3850 is resolved,
// remove this and just use templated getLocalView<Device>(ReadOnly).
template<typename MV>
struct GetReadOnly
{
  using DevView = typename MV::dual_view_type::t_dev::const_type;
  using HostView = typename MV::dual_view_type::t_host::const_type;

  template<typename exec_space>
  static DevView get(const MV& x, typename std::enable_if<std::is_same<exec_space, typename MV::execution_space>::value>::type* = nullptr)
  {
    return x.getLocalViewDevice(Tpetra::Access::ReadOnly);
  }

  template<typename exec_space>
  static HostView get(const MV& x, typename std::enable_if<!std::is_same<exec_space, typename MV::execution_space>::value>::type* = nullptr)
  {
    return x.getLocalViewHost(Tpetra::Access::ReadOnly);
  }
};


template<class MV, class ResultView, bool runOnDevice>
void idotLocal(const ResultView& localResult,
               const MV& X,
               const MV& Y)
{
  using pair_type = Kokkos::pair<size_t, size_t>;
  using exec_space = typename std::conditional<runOnDevice, typename MV::execution_space, Kokkos::DefaultHostExecutionSpace>::type;
  //if the execution space can access localResult, use it directly. Otherwise need to make a copy.
  static_assert(Kokkos::SpaceAccessibility<exec_space, typename ResultView::memory_space>::accessible,
      "idotLocal: Execution space must be able to access localResult");
  //If Dot executes on Serial, it requires the result to be HostSpace. If localResult is CudaUVMSpace,
  //we can just treat it like HostSpace.
  Kokkos::View<typename ResultView::data_type, typename exec_space::memory_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
    localResultUnmanaged(localResult.data(), localResult.extent(0));
  const size_t numRows = X.getLocalLength ();
  const pair_type rowRange (0, numRows);
  const size_t X_numVecs = X.getNumVectors ();
  const size_t Y_numVecs = Y.getNumVectors ();
  const size_t numVecs = X_numVecs > Y_numVecs ? X_numVecs : Y_numVecs;
  // Check compatibility of number of columns; allow special cases of
  // a multivector dot a single vector, or vice versa.
  if (X_numVecs != Y_numVecs &&
      X_numVecs != size_t (1) &&
      Y_numVecs != size_t (1)) {
    std::ostringstream os;
    os << "Tpetra::idot: X.getNumVectors() = " << X_numVecs
       << " != Y.getNumVectors() = " << Y_numVecs
       << ", but neither is 1.";
    throw std::invalid_argument (os.str ());
  }
  auto X_lcl = GetReadOnly<MV>::template get<exec_space>(X);
  auto Y_lcl = GetReadOnly<MV>::template get<exec_space>(Y);
  //Can the multivector dot kernel be used?
  bool useMVDot = X.isConstantStride() && Y.isConstantStride() && X_numVecs == Y_numVecs;
  if(useMVDot)
  {
    if (numVecs == size_t (1)) {
      auto X_lcl_1d = Kokkos::subview (X_lcl, rowRange, 0);
      auto Y_lcl_1d = Kokkos::subview (Y_lcl, rowRange, 0);
      auto result_0d = Kokkos::subview (localResultUnmanaged, 0);
      KokkosBlas::dot (result_0d, X_lcl_1d, Y_lcl_1d);
    }
    else {
      auto X_lcl_2d = Kokkos::subview (X_lcl, rowRange, pair_type (0, X_numVecs));
      auto Y_lcl_2d = Kokkos::subview (Y_lcl, rowRange, pair_type (0, Y_numVecs));
      KokkosBlas::dot (localResultUnmanaged, X_lcl_2d, Y_lcl_2d);
    }
  }
  else
  {
    auto XWhichVectors = Tpetra::getMultiVectorWhichVectors(X);
    auto YWhichVectors = Tpetra::getMultiVectorWhichVectors(Y);
    //Need to compute each dot individually, using 1D subviews of X_lcl and Y_lcl
    for(size_t vec = 0; vec < numVecs; vec++) {
      //Get the Vectors for each column of X and Y that will be dotted together.
      //Note: "which vectors" is not populated for constant stride MVs (but it's the identity mapping)
      size_t Xj = (numVecs == X_numVecs) ? vec : 0;
      Xj = X.isConstantStride() ? Xj : XWhichVectors[Xj];
      size_t Yj = (numVecs == Y_numVecs) ? vec : 0;
      Yj = Y.isConstantStride() ? Yj : YWhichVectors[Yj];
      auto Xcol = Kokkos::subview(X_lcl, rowRange, Xj);
      auto Ycol = Kokkos::subview(Y_lcl, rowRange, Yj);

      //Compute the rank-1 dot product, and place the result in an element of localResult
      KokkosBlas::dot(Kokkos::subview(localResultUnmanaged, vec), Xcol, Ycol);
    }
  }
}

//Helper to avoid extra instantiations of KokkosBlas::dot and iallreduce.
template<typename MV, typename ResultView>
struct IdotHelper
{
  using dot_type = typename MV::dot_type;

  //First version: use result directly
  template<typename exec_space>
  static std::shared_ptr< ::Tpetra::Details::CommRequest> run(
      const ResultView& globalResult, const MV& X, const MV& Y,
      typename std::enable_if<Kokkos::SpaceAccessibility<exec_space, typename ResultView::memory_space>::accessible>::type* = nullptr)

  {
    constexpr bool runOnDevice = std::is_same<exec_space, typename MV::execution_space>::value;
    idotLocal<MV, ResultView, runOnDevice>(globalResult, X, Y);
    //Fence because we're giving result of device kernel to MPI
    if(runOnDevice)
      exec_space().fence();
    auto comm = X.getMap()->getComm();
    return iallreduce(globalResult, globalResult, ::Teuchos::REDUCE_SUM, *comm);
  }

  //Second version: use a temporary local result, because exec_space can't write to globalResult
  template<typename exec_space>
  static std::shared_ptr< ::Tpetra::Details::CommRequest> run(
      const ResultView& globalResult, const MV& X, const MV& Y,
      typename std::enable_if<!Kokkos::SpaceAccessibility<exec_space, typename ResultView::memory_space>::accessible>::type* = nullptr)
  {
    constexpr bool runOnDevice = std::is_same<exec_space, typename MV::execution_space>::value;
    Kokkos::View<dot_type*, typename exec_space::memory_space> localResult(Kokkos::ViewAllocateWithoutInitializing("idot:localResult"), X.getNumVectors());
    idotLocal<MV, decltype(localResult), runOnDevice>(localResult, X, Y);
    //Fence because we're giving result of device kernel to MPI
    if(runOnDevice)
      exec_space().fence();
    auto comm = X.getMap()->getComm();
    return iallreduce(localResult, globalResult, ::Teuchos::REDUCE_SUM, *comm);
  }

};


template<class MV, class ResultView>
std::shared_ptr< ::Tpetra::Details::CommRequest>
idotImpl(const ResultView& globalResult,
         const MV& X,
         const MV& Y)
{
  static_assert(std::is_same<typename ResultView::non_const_value_type, typename MV::dot_type>::value,
      "Tpetra::idot: result view's element type must match MV::dot_type");

  // Execution space to use for dot kernel(s) is whichever has access to up-to-date X.
  if(X.need_sync_device())
  {
    //run on host.
    return IdotHelper<MV, ResultView>::template run<Kokkos::DefaultHostExecutionSpace>(globalResult, X, Y);
  }
  else
  {
    //run on device.
    return IdotHelper<MV, ResultView>::template run<typename MV::execution_space>(globalResult, X, Y);
  }
}
} // namespace Details

//
// SKIP DOWN TO HERE
//

template<class SC, class LO, class GO, class NT>
std::shared_ptr< ::Tpetra::Details::CommRequest>
idot (typename ::Tpetra::MultiVector<SC, LO, GO, NT>::dot_type* resultRaw,
      const ::Tpetra::MultiVector<SC, LO, GO, NT>& X,
      const ::Tpetra::MultiVector<SC, LO, GO, NT>& Y)
{
  using dot_type = typename ::Tpetra::Vector<SC, LO, GO, NT>::dot_type;
  const size_t X_numVecs = X.getNumVectors ();
  const size_t Y_numVecs = Y.getNumVectors ();
  const size_t numVecs = (X_numVecs > Y_numVecs) ? X_numVecs : Y_numVecs;
  Kokkos::View<dot_type*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>
    resultView(resultRaw, numVecs);
  return Details::idotImpl(resultView, X, Y);
}

template<class SC, class LO, class GO, class NT>
std::shared_ptr< ::Tpetra::Details::CommRequest>
idot (const Kokkos::View<typename ::Tpetra::MultiVector<SC, LO, GO, NT>::dot_type*,
        typename ::Tpetra::MultiVector<SC, LO, GO, NT>::device_type>& result,
      const ::Tpetra::MultiVector<SC, LO, GO, NT>& X,
      const ::Tpetra::MultiVector<SC, LO, GO, NT>& Y)
{
  return Details::idotImpl(result, X, Y);
}

template<class SC, class LO, class GO, class NT>
std::shared_ptr< ::Tpetra::Details::CommRequest>
idot (const Kokkos::View<typename ::Tpetra::Vector<SC, LO, GO, NT>::dot_type,
        typename ::Tpetra::Vector<SC, LO, GO, NT>::device_type>& result,
      const ::Tpetra::Vector<SC, LO, GO, NT>& X,
      const ::Tpetra::Vector<SC, LO, GO, NT>& Y)
{
  using dot_type = typename ::Tpetra::Vector<SC, LO, GO, NT>::dot_type;
  using result_device_t = typename ::Tpetra::Vector<SC, LO, GO, NT>::device_type;
  Kokkos::View<dot_type*, result_device_t, Kokkos::MemoryTraits<Kokkos::Unmanaged>> result1D(result.data(), 1);
  return Details::idotImpl(result1D, X, Y);
}

} // namespace Tpetra

#endif // TPETRA_DETAILS_IDOT_HPP