Tpetra_Details_StaticView.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_STATICVIEW_HPP
#define TPETRA_DETAILS_STATICVIEW_HPP

#include "TpetraCore_config.h"
//#include "Tpetra_Details_Behavior.hpp"
#include "Kokkos_DualView.hpp"

namespace Tpetra {
namespace Details {
namespace Impl {

template<class MemorySpace>
class StaticKokkosAllocation {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  // Allocation automatically registers deallocation to happen at
  // Kokkos::finalize.  Reallocation only happens if needed.
  static void* resize (MemorySpace space, const size_t size);
};

#ifdef KOKKOS_ENABLE_CUDA
template<>
class StaticKokkosAllocation<Kokkos::CudaSpace> {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  static void* resize (Kokkos::CudaSpace space, const size_t size);
};

template<>
class StaticKokkosAllocation<Kokkos::CudaUVMSpace> {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  static void* resize (Kokkos::CudaUVMSpace space, const size_t size);
};

template<>
class StaticKokkosAllocation<Kokkos::CudaHostPinnedSpace> {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  static void* resize (Kokkos::CudaHostPinnedSpace space, const size_t size);
};
#endif // KOKKOS_ENABLE_CUDA

#ifdef KOKKOS_ENABLE_HIP
template<>
class StaticKokkosAllocation<Kokkos::HIPSpace> {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  static void* resize (Kokkos::HIPSpace space, const size_t size);
};

template<>
class StaticKokkosAllocation<Kokkos::HIPHostPinnedSpace> {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  static void* resize (Kokkos::HIPHostPinnedSpace space, const size_t size);
};
#endif // KOKKOS_ENABLE_HIP

template<>
class StaticKokkosAllocation<Kokkos::HostSpace> {
public:
  StaticKokkosAllocation () = delete;
  ~StaticKokkosAllocation () = delete;
  StaticKokkosAllocation (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation& operator= (const StaticKokkosAllocation&) = delete;
  StaticKokkosAllocation (StaticKokkosAllocation&&) = delete;
  StaticKokkosAllocation& operator= (StaticKokkosAllocation&&) = delete;

  static void* resize (Kokkos::HostSpace space, const size_t size);
};

template<class ValueType, class MemorySpace>
ValueType*
getStaticKokkosMemory (MemorySpace space,
                       const size_t num_entries,
                       const size_t value_size = sizeof (ValueType))
{
  void* ptr = StaticKokkosAllocation<MemorySpace>::resize
    (space, num_entries * value_size);
  return reinterpret_cast<ValueType*> (ptr);
}

} // namespace Impl

template<class ValueType, class DeviceType>
Kokkos::View<ValueType*, DeviceType>
getStatic1dView (const size_t size)
{
  using Impl::getStaticKokkosMemory;
  using mem_space = typename DeviceType::memory_space;
  using view_type = Kokkos::View<ValueType*, DeviceType>;

  ValueType* ptr = getStaticKokkosMemory<ValueType> (mem_space (), size);
  return view_type (ptr, size);
}

template<class ValueType, class DeviceType>
Kokkos::View<ValueType**, Kokkos::LayoutLeft, DeviceType>
getStatic2dView (const size_t num_rows, const size_t num_cols)
{
  using Impl::getStaticKokkosMemory;
  using mem_space = typename DeviceType::memory_space;
  using view_type = Kokkos::View<ValueType**, Kokkos::LayoutLeft, DeviceType>;

  const size_t size = num_rows * num_cols;
  ValueType* ptr = getStaticKokkosMemory<ValueType> (mem_space (), size);
  return view_type (ptr, num_rows, num_cols);
}

template<class ValueType, class DeviceType>
Kokkos::DualView<ValueType**, Kokkos::LayoutLeft, DeviceType>
getStatic2dDualView (const size_t num_rows, const size_t num_cols)
{
  using dual_view_type =
    Kokkos::DualView<ValueType**, Kokkos::LayoutLeft, DeviceType>;
  using d_view_type = typename dual_view_type::t_dev;
  using h_view_type = typename dual_view_type::t_host;

  auto d_view = getStatic2dView<ValueType, DeviceType> (num_rows, num_cols);
  // Preserve the invariant that Kokkos::create_mirror_view returns
  // the input View here, if and only if it would have returned it if
  // the allocating View constructor were called.
  h_view_type h_view;
  if (std::is_same<typename d_view_type::memory_space,
                   typename h_view_type::memory_space>::value) {
    h_view = Kokkos::create_mirror_view (d_view);
  }
  else {
    h_view = getStatic2dView<ValueType,
      typename h_view_type::device_type> (num_rows, num_cols);
  }

  return dual_view_type (d_view, h_view);
}


} // namespace Details
} // namespace Tpetra

#endif // TPETRA_DETAILS_STATICVIEW_HPP