Kokkos_DynamicView.hpp

/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Christian R. Trott (crtrott@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#ifndef KOKKOS_DYNAMIC_VIEW_HPP
#define KOKKOS_DYNAMIC_VIEW_HPP

#include <cstdio>

#include <Kokkos_Core.hpp>
#include <impl/Kokkos_Error.hpp>

namespace Kokkos {
namespace Experimental {

// Simple metafunction for choosing memory space
// In the current implementation, if memory_space == CudaSpace,
// use CudaUVMSpace for the chunk 'array' allocation, which
// contains will contain pointers to chunks of memory allocated
// in CudaSpace
namespace Impl {
template <class MemSpace>
struct ChunkArraySpace {
  using memory_space = MemSpace;
};

#ifdef KOKKOS_ENABLE_CUDA
template <>
struct ChunkArraySpace<Kokkos::CudaSpace> {
  using memory_space = typename Kokkos::CudaUVMSpace;
};
#endif
#ifdef KOKKOS_ENABLE_HIP
template <>
struct ChunkArraySpace<Kokkos::Experimental::HIPSpace> {
  using memory_space = typename Kokkos::Experimental::HIPHostPinnedSpace;
};
#endif
}  // end namespace Impl

template <typename DataType, typename... P>
class DynamicView : public Kokkos::ViewTraits<DataType, P...> {
 public:
  typedef Kokkos::ViewTraits<DataType, P...> traits;

 private:
  template <class, class...>
  friend class DynamicView;

  typedef Kokkos::Impl::SharedAllocationTracker track_type;

  static_assert(traits::rank == 1 && traits::rank_dynamic == 1,
                "DynamicView must be rank-one");

  // It is assumed that the value_type is trivially copyable;
  // when this is not the case, potential problems can occur.
  static_assert(std::is_same<typename traits::specialize, void>::value,
                "DynamicView only implemented for non-specialized View type");

  template <class Space, bool = Kokkos::Impl::MemorySpaceAccess<
                             Space, typename traits::memory_space>::accessible>
  struct verify_space {
    KOKKOS_FORCEINLINE_FUNCTION static void check() {}
  };

  template <class Space>
  struct verify_space<Space, false> {
    KOKKOS_FORCEINLINE_FUNCTION static void check() {
      Kokkos::abort(
          "Kokkos::DynamicView ERROR: attempt to access inaccessible memory "
          "space");
    };
  };

 private:
  track_type m_track;
  typename traits::value_type**
      m_chunks;            // array of pointers to 'chunks' of memory
  unsigned m_chunk_shift;  // ceil(log2(m_chunk_size))
  unsigned m_chunk_mask;   // m_chunk_size - 1
  unsigned m_chunk_max;  // number of entries in the chunk array - each pointing
                         // to a chunk of extent == m_chunk_size entries
  unsigned m_chunk_size;  // 2 << (m_chunk_shift - 1)

 public:
  //----------------------------------------------------------------------

  typedef DynamicView<typename traits::data_type, typename traits::device_type>
      array_type;

  typedef DynamicView<typename traits::const_data_type,
                      typename traits::device_type>
      const_type;

  typedef DynamicView<typename traits::non_const_data_type,
                      typename traits::device_type>
      non_const_type;

  typedef DynamicView HostMirror;

  typedef Kokkos::Device<typename traits::device_type::execution_space,
                         Kokkos::AnonymousSpace>
      uniform_device;
  typedef array_type uniform_type;
  typedef const_type uniform_const_type;
  typedef array_type uniform_runtime_type;
  typedef const_type uniform_runtime_const_type;
  typedef DynamicView<typename traits::data_type, uniform_device>
      uniform_nomemspace_type;
  typedef DynamicView<typename traits::const_data_type, uniform_device>
      uniform_const_nomemspace_type;
  typedef DynamicView<typename traits::data_type, uniform_device>
      uniform_runtime_nomemspace_type;
  typedef DynamicView<typename traits::const_data_type, uniform_device>
      uniform_runtime_const_nomemspace_type;

  //----------------------------------------------------------------------

  enum { Rank = 1 };

  KOKKOS_INLINE_FUNCTION
  size_t allocation_extent() const noexcept {
    uintptr_t n = *reinterpret_cast<const uintptr_t*>(m_chunks + m_chunk_max);
    return (n << m_chunk_shift);
  }

  KOKKOS_INLINE_FUNCTION
  size_t chunk_size() const noexcept { return m_chunk_size; }

  KOKKOS_INLINE_FUNCTION
  size_t size() const noexcept {
    size_t extent_0 =
        *reinterpret_cast<const size_t*>(m_chunks + m_chunk_max + 1);
    return extent_0;
  }

  template <typename iType>
  KOKKOS_INLINE_FUNCTION size_t extent(const iType& r) const {
    return r == 0 ? size() : 1;
  }

  template <typename iType>
  KOKKOS_INLINE_FUNCTION size_t extent_int(const iType& r) const {
    return r == 0 ? size() : 1;
  }

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
  KOKKOS_INLINE_FUNCTION size_t dimension_0() const { return size(); }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_1() const { return 1; }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_2() const { return 1; }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_3() const { return 1; }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_4() const { return 1; }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_5() const { return 1; }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_6() const { return 1; }
  KOKKOS_INLINE_FUNCTION constexpr size_t dimension_7() const { return 1; }
#endif

  KOKKOS_INLINE_FUNCTION constexpr size_t stride_0() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_1() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_2() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_3() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_4() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_5() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_6() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr size_t stride_7() const { return 0; }

  template <typename iType>
  KOKKOS_INLINE_FUNCTION void stride(iType* const s) const {
    *s = 0;
  }

  //----------------------------------------
  // Allocation tracking properties

  KOKKOS_INLINE_FUNCTION
  int use_count() const { return m_track.use_count(); }

  inline const std::string label() const {
    return m_track.template get_label<typename traits::memory_space>();
  }

  //----------------------------------------------------------------------
  // Range span is the span which contains all members.

  typedef typename traits::value_type& reference_type;
  typedef typename traits::value_type* pointer_type;

  enum {
    reference_type_is_lvalue_reference =
        std::is_lvalue_reference<reference_type>::value
  };

  KOKKOS_INLINE_FUNCTION constexpr bool span_is_contiguous() const {
    return false;
  }
  KOKKOS_INLINE_FUNCTION constexpr size_t span() const { return 0; }
  KOKKOS_INLINE_FUNCTION constexpr pointer_type data() const { return 0; }

  //----------------------------------------

  template <typename I0, class... Args>
  KOKKOS_INLINE_FUNCTION reference_type
  operator()(const I0& i0, const Args&... /*args*/) const {
    static_assert(Kokkos::Impl::are_integral<I0, Args...>::value,
                  "Indices must be integral type");

    DynamicView::template verify_space<
        Kokkos::Impl::ActiveExecutionMemorySpace>::check();

    // Which chunk is being indexed.
    const uintptr_t ic = uintptr_t(i0 >> m_chunk_shift);

    typename traits::value_type* volatile* const ch = m_chunks + ic;

    // Do bounds checking if enabled or if the chunk pointer is zero.
    // If not bounds checking then we assume a non-zero pointer is valid.

#if !defined(KOKKOS_ENABLE_DEBUG_BOUNDS_CHECK)
    if (nullptr == *ch)
#endif
    {
      // Verify that allocation of the requested chunk in in progress.

      // The allocated chunk counter is m_chunks[ m_chunk_max ]
      const uintptr_t n =
          *reinterpret_cast<uintptr_t volatile*>(m_chunks + m_chunk_max);

      if (n <= ic) {
        Kokkos::abort("Kokkos::DynamicView array bounds error");
      }

      // Allocation of this chunk is in progress
      // so wait for allocation to complete.
      while (nullptr == *ch)
        ;
    }

    return (*ch)[i0 & m_chunk_mask];
  }

  //----------------------------------------
  template <typename IntType>
  inline typename std::enable_if<
      std::is_integral<IntType>::value &&
      Kokkos::Impl::MemorySpaceAccess<
          Kokkos::HostSpace,
          typename Impl::ChunkArraySpace<
              typename traits::memory_space>::memory_space>::accessible>::type
  resize_serial(IntType const& n) {
    typedef typename traits::value_type local_value_type;
    typedef local_value_type* value_pointer_type;

    const uintptr_t NC =
        (n + m_chunk_mask) >>
        m_chunk_shift;  // New total number of chunks needed for resize

    if (m_chunk_max < NC) {
      Kokkos::abort("DynamicView::resize_serial exceeded maximum size");
    }

    // *m_chunks[m_chunk_max] stores the current number of chunks being used
    uintptr_t* const pc = reinterpret_cast<uintptr_t*>(m_chunks + m_chunk_max);

    if (*pc < NC) {
      while (*pc < NC) {
        m_chunks[*pc] = reinterpret_cast<value_pointer_type>(
            typename traits::memory_space().allocate(sizeof(local_value_type)
                                                     << m_chunk_shift));
        ++*pc;
      }
    } else {
      while (NC + 1 <= *pc) {
        --*pc;
        typename traits::memory_space().deallocate(
            m_chunks[*pc], sizeof(local_value_type) << m_chunk_shift);
        m_chunks[*pc] = nullptr;
      }
    }
    // *m_chunks[m_chunk_max+1] stores the 'extent' requested by resize
    *(pc + 1) = n;
  }

  //----------------------------------------------------------------------

  ~DynamicView()                  = default;
  DynamicView()                   = default;
  DynamicView(DynamicView&&)      = default;
  DynamicView(const DynamicView&) = default;
  DynamicView& operator=(DynamicView&&) = default;
  DynamicView& operator=(const DynamicView&) = default;

  template <class RT, class... RP>
  DynamicView(const DynamicView<RT, RP...>& rhs)
      : m_track(rhs.m_track),
        m_chunks((typename traits::value_type**)rhs.m_chunks),
        m_chunk_shift(rhs.m_chunk_shift),
        m_chunk_mask(rhs.m_chunk_mask),
        m_chunk_max(rhs.m_chunk_max),
        m_chunk_size(rhs.m_chunk_size) {
    typedef typename DynamicView<RT, RP...>::traits SrcTraits;
    typedef Kokkos::Impl::ViewMapping<traits, SrcTraits, void> Mapping;
    static_assert(Mapping::is_assignable,
                  "Incompatible DynamicView copy construction");
  }

  //----------------------------------------------------------------------

  struct Destroy {
    typename traits::value_type** m_chunks;
    unsigned m_chunk_max;
    bool m_destroy;
    unsigned m_chunk_size;

    // Initialize or destroy array of chunk pointers.
    // Two entries beyond the max chunks are allocation counters.
    inline void operator()(unsigned i) const {
      if (m_destroy && i < m_chunk_max && nullptr != m_chunks[i]) {
        typename traits::memory_space().deallocate(m_chunks[i], m_chunk_size);
      }
      m_chunks[i] = nullptr;
    }

    void execute(bool arg_destroy) {
      typedef Kokkos::RangePolicy<typename HostSpace::execution_space> Range;
      // typedef Kokkos::RangePolicy< typename Impl::ChunkArraySpace< typename
      // traits::memory_space >::memory_space::execution_space > Range ;

      m_destroy = arg_destroy;

      Kokkos::Impl::ParallelFor<Destroy, Range> closure(
          *this,
          Range(0, m_chunk_max + 2));  // Add 2 to 'destroy' extra slots storing
                                       // num_chunks and extent; previously + 1

      closure.execute();

      typename traits::execution_space().fence();
      // Impl::ChunkArraySpace< typename traits::memory_space
      // >::memory_space::execution_space().fence();
    }

    void construct_shared_allocation() { execute(false); }

    void destroy_shared_allocation() { execute(true); }

    Destroy()               = default;
    Destroy(Destroy&&)      = default;
    Destroy(const Destroy&) = default;
    Destroy& operator=(Destroy&&) = default;
    Destroy& operator=(const Destroy&) = default;

    Destroy(typename traits::value_type** arg_chunk,
            const unsigned arg_chunk_max, const unsigned arg_chunk_size)
        : m_chunks(arg_chunk),
          m_chunk_max(arg_chunk_max),
          m_destroy(false),
          m_chunk_size(arg_chunk_size) {}
  };

  explicit inline DynamicView(const std::string& arg_label,
                              const unsigned min_chunk_size,
                              const unsigned max_extent)
      : m_track(),
        m_chunks(nullptr)
        // The chunk size is guaranteed to be a power of two
        ,
        m_chunk_shift(Kokkos::Impl::integral_power_of_two_that_contains(
            min_chunk_size))  // div ceil(log2(min_chunk_size))
        ,
        m_chunk_mask((1 << m_chunk_shift) - 1)  // mod
        ,
        m_chunk_max((max_extent + m_chunk_mask) >>
                    m_chunk_shift)  // max num pointers-to-chunks in array
        ,
        m_chunk_size(2 << (m_chunk_shift - 1)) {
    typedef typename Impl::ChunkArraySpace<
        typename traits::memory_space>::memory_space chunk_array_memory_space;
    // A functor to deallocate all of the chunks upon final destruction
    typedef Kokkos::Impl::SharedAllocationRecord<chunk_array_memory_space,
                                                 Destroy>
        record_type;

    // Allocate chunk pointers and allocation counter
    record_type* const record =
        record_type::allocate(chunk_array_memory_space(), arg_label,
                              (sizeof(pointer_type) * (m_chunk_max + 2)));
    // Allocate + 2 extra slots so that *m_chunk[m_chunk_max] ==
    // num_chunks_alloc and *m_chunk[m_chunk_max+1] == extent This must match in
    // Destroy's execute(...) method

    m_chunks = reinterpret_cast<pointer_type*>(record->data());

    record->m_destroy = Destroy(m_chunks, m_chunk_max, m_chunk_size);

    // Initialize to zero
    record->m_destroy.construct_shared_allocation();

    m_track.assign_allocated_record_to_uninitialized(record);
  }
};

}  // namespace Experimental
}  // namespace Kokkos

namespace Kokkos {

template <class T, class... P>
inline typename Kokkos::Experimental::DynamicView<T, P...>::HostMirror
create_mirror_view(const Kokkos::Experimental::DynamicView<T, P...>& src) {
  return src;
}

template <class T, class... DP, class... SP>
inline void deep_copy(const View<T, DP...>& dst,
                      const Kokkos::Experimental::DynamicView<T, SP...>& src) {
  typedef View<T, DP...> dst_type;
  typedef Kokkos::Experimental::DynamicView<T, SP...> src_type;

  typedef typename ViewTraits<T, DP...>::execution_space dst_execution_space;
  typedef typename ViewTraits<T, SP...>::memory_space src_memory_space;

  enum {
    DstExecCanAccessSrc =
        Kokkos::Impl::SpaceAccessibility<dst_execution_space,
                                         src_memory_space>::accessible
  };

  if (DstExecCanAccessSrc) {
    // Copying data between views in accessible memory spaces and either
    // non-contiguous or incompatible shape.
    Kokkos::Impl::ViewRemap<dst_type, src_type>(dst, src);
  } else {
    Kokkos::Impl::throw_runtime_exception(
        "deep_copy given views that would require a temporary allocation");
  }
}

template <class T, class... DP, class... SP>
inline void deep_copy(const Kokkos::Experimental::DynamicView<T, DP...>& dst,
                      const View<T, SP...>& src) {
  typedef Kokkos::Experimental::DynamicView<T, SP...> dst_type;
  typedef View<T, DP...> src_type;

  typedef typename ViewTraits<T, DP...>::execution_space dst_execution_space;
  typedef typename ViewTraits<T, SP...>::memory_space src_memory_space;

  enum {
    DstExecCanAccessSrc =
        Kokkos::Impl::SpaceAccessibility<dst_execution_space,
                                         src_memory_space>::accessible
  };

  if (DstExecCanAccessSrc) {
    // Copying data between views in accessible memory spaces and either
    // non-contiguous or incompatible shape.
    Kokkos::Impl::ViewRemap<dst_type, src_type>(dst, src);
  } else {
    Kokkos::Impl::throw_runtime_exception(
        "deep_copy given views that would require a temporary allocation");
  }
}

namespace Impl {
template <class Arg0, class... DP, class... SP>
struct CommonSubview<Kokkos::Experimental::DynamicView<DP...>,
                     Kokkos::Experimental::DynamicView<SP...>, 1, Arg0> {
  typedef Kokkos::Experimental::DynamicView<DP...> DstType;
  typedef Kokkos::Experimental::DynamicView<SP...> SrcType;
  typedef DstType dst_subview_type;
  typedef SrcType src_subview_type;
  dst_subview_type dst_sub;
  src_subview_type src_sub;
  CommonSubview(const DstType& dst, const SrcType& src, const Arg0& /*arg0*/)
      : dst_sub(dst), src_sub(src) {}
};

template <class... DP, class SrcType, class Arg0>
struct CommonSubview<Kokkos::Experimental::DynamicView<DP...>, SrcType, 1,
                     Arg0> {
  typedef Kokkos::Experimental::DynamicView<DP...> DstType;
  typedef DstType dst_subview_type;
  typedef typename Kokkos::Subview<SrcType, Arg0> src_subview_type;
  dst_subview_type dst_sub;
  src_subview_type src_sub;
  CommonSubview(const DstType& dst, const SrcType& src, const Arg0& arg0)
      : dst_sub(dst), src_sub(src, arg0) {}
};

template <class DstType, class... SP, class Arg0>
struct CommonSubview<DstType, Kokkos::Experimental::DynamicView<SP...>, 1,
                     Arg0> {
  typedef Kokkos::Experimental::DynamicView<SP...> SrcType;
  typedef typename Kokkos::Subview<DstType, Arg0> dst_subview_type;
  typedef SrcType src_subview_type;
  dst_subview_type dst_sub;
  src_subview_type src_sub;
  CommonSubview(const DstType& dst, const SrcType& src, const Arg0& arg0)
      : dst_sub(dst, arg0), src_sub(src) {}
};

template <class... DP, class ViewTypeB, class Layout, class ExecSpace,
          typename iType>
struct ViewCopy<Kokkos::Experimental::DynamicView<DP...>, ViewTypeB, Layout,
                ExecSpace, 1, iType, false> {
  Kokkos::Experimental::DynamicView<DP...> a;
  ViewTypeB b;

  typedef Kokkos::RangePolicy<ExecSpace, Kokkos::IndexType<iType>> policy_type;

  ViewCopy(const Kokkos::Experimental::DynamicView<DP...>& a_,
           const ViewTypeB& b_)
      : a(a_), b(b_) {
    Kokkos::parallel_for("Kokkos::ViewCopy-1D", policy_type(0, b.extent(0)),
                         *this);
  }

  KOKKOS_INLINE_FUNCTION
  void operator()(const iType& i0) const { a(i0) = b(i0); };
};

template <class... DP, class... SP, class Layout, class ExecSpace,
          typename iType>
struct ViewCopy<Kokkos::Experimental::DynamicView<DP...>,
                Kokkos::Experimental::DynamicView<SP...>, Layout, ExecSpace, 1,
                iType, false> {
  Kokkos::Experimental::DynamicView<DP...> a;
  Kokkos::Experimental::DynamicView<SP...> b;

  typedef Kokkos::RangePolicy<ExecSpace, Kokkos::IndexType<iType>> policy_type;

  ViewCopy(const Kokkos::Experimental::DynamicView<DP...>& a_,
           const Kokkos::Experimental::DynamicView<SP...>& b_)
      : a(a_), b(b_) {
    const iType n = std::min(a.extent(0), b.extent(0));
    Kokkos::parallel_for("Kokkos::ViewCopy-1D", policy_type(0, n), *this);
  }

  KOKKOS_INLINE_FUNCTION
  void operator()(const iType& i0) const { a(i0) = b(i0); };
};

}  // namespace Impl
}  // namespace Kokkos

#endif /* #ifndef KOKKOS_DYNAMIC_VIEW_HPP */