Tpetra_Details_temporaryViewUtils.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_TEMPVIEWUTILS_HPP
#define TPETRA_DETAILS_TEMPVIEWUTILS_HPP

#include "Kokkos_Core.hpp"
#include "Tpetra_Details_isInterComm.hpp"

namespace Tpetra
{
namespace Details
{
namespace TempView
{

template<typename MemorySpace>
struct AlwaysMPISafe
{
  enum : bool {value = false};
};

template<>
struct AlwaysMPISafe<Kokkos::HostSpace>
{
  enum : bool {value = true};
};

#ifdef KOKKOS_ENABLE_CUDA
template<>
struct AlwaysMPISafe<Kokkos::CudaHostPinnedSpace>
{
  enum : bool {value = true};
};
#endif

#ifdef KOKKOS_ENABLE_HIP
template<>
struct AlwaysMPISafe<Kokkos::HIPHostPinnedSpace>
{
  enum : bool {value = true};
};
#endif

template<typename View1, typename View2>
struct UnifiedContiguousLayout
{
  using L1 = typename View1::array_layout;
  using L2 = typename View2::array_layout;
  enum : bool {EitherLeft = std::is_same<L1, Kokkos::LayoutLeft>::value || std::is_same<L2, Kokkos::LayoutLeft>::value};
enum : bool {BothStride = std::is_same<L1, Kokkos::LayoutStride>::value && std::is_same<L2, Kokkos::LayoutStride>::value};
  using type = typename std::conditional<EitherLeft || BothStride, Kokkos::LayoutLeft, Kokkos::LayoutRight>::type;
};

template<typename SrcView, typename Layout, typename std::enable_if<!std::is_same<typename SrcView::array_layout, Layout>::value>::type* = nullptr>
Kokkos::View<typename SrcView::data_type, Layout, typename SrcView::device_type>
toLayout(const SrcView& src)
{
  static_assert(!std::is_same<Kokkos::LayoutStride, Layout>::value,
      "TempView::toLayout: Layout must be contiguous (not LayoutStride)");
  Layout layout(src.extent(0), src.extent(1));
  Kokkos::View<typename SrcView::non_const_data_type, Layout, typename SrcView::device_type> dst(Kokkos::ViewAllocateWithoutInitializing(src.label()), layout);
  Kokkos::deep_copy(dst, src);
  return dst;
}

template<typename SrcView, typename Layout, typename std::enable_if<std::is_same<typename SrcView::array_layout, Layout>::value>::type* = nullptr>
Kokkos::View<typename SrcView::data_type, Layout, typename SrcView::device_type>
toLayout(const SrcView& src)
{
  if(src.span_is_contiguous())
  {
    return src;
  }
  else
  {
    //Even though the layout is already correct, it's not contiguous.
    Layout layout(src.extent(0), src.extent(1));
    Kokkos::View<typename SrcView::non_const_data_type, Layout, typename SrcView::device_type>
      result(Kokkos::ViewAllocateWithoutInitializing(src.label()), layout);
    Kokkos::deep_copy(result, src);
    return result;
  }
}

template<typename SrcView, bool AssumeGPUAware, typename = typename std::enable_if<AssumeGPUAware || AlwaysMPISafe<typename SrcView::memory_space>::value>::type>
SrcView
toMPISafe(const SrcView& src)
{
  using SrcLayout = typename SrcView::array_layout;
  static_assert(!std::is_same<SrcLayout, Kokkos::LayoutStride>::value, "toMPISafe requires that SrcView is contiguous");
  return toLayout<SrcView, SrcLayout>(src);
}

template<typename SrcView, bool AssumeGPUAware, typename = typename std::enable_if<!(AssumeGPUAware || AlwaysMPISafe<typename SrcView::memory_space>::value)>::type>
decltype(Kokkos::create_mirror_view_and_copy(std::declval<Kokkos::HostSpace>(), std::declval<SrcView>()))
toMPISafe(const SrcView& src)
{
  using SrcLayout = typename SrcView::array_layout;
  static_assert(!std::is_same<SrcLayout, Kokkos::LayoutStride>::value, "toMPISafe requires that SrcView is contiguous");
  auto srcContig = toLayout<SrcView, SrcLayout>(src);
  return Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), srcContig);
}

}}} //namespace Tpetra::Details::TempView

#endif