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