Tpetra_Details_DualViewUtil.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_DUALVIEWUTIL_HPP
#define TPETRA_DETAILS_DUALVIEWUTIL_HPP

#include "TpetraCore_config.h"
#include "Kokkos_DualView.hpp"
#include "Teuchos_ArrayView.hpp"
#include <ostream>
#include <string>

namespace Tpetra {

namespace Details {

auto view_alloc_no_init(const std::string& label) -> decltype(Kokkos::view_alloc(label, Kokkos::WithoutInitializing));

template <class ElementType, class DeviceType>
void makeDualViewFromOwningHostView(Kokkos::DualView<ElementType*, DeviceType>& dv,
                                    const typename Kokkos::DualView<ElementType*, DeviceType>::t_host& hostView) {
  using execution_space = typename DeviceType::execution_space;
  using dual_view_type  = Kokkos::DualView<ElementType*, DeviceType>;

  if constexpr (Kokkos::SpaceAccessibility<Kokkos::HostSpace, typename DeviceType::memory_space>::accessible) {
    // DualView only references one View, so we pass in the same View twice
    dv = dual_view_type(hostView, hostView);
  } else {
    typename Kokkos::DualView<ElementType*, DeviceType>::t_dev devView;
    if (dv.extent(0) == hostView.extent(0))
      devView = dv.view_device();
    else
      devView = Kokkos::create_mirror_view(DeviceType(), hostView);
    // DEEP_COPY REVIEW - DEVICE-TO-HOSTMIRROR
    Kokkos::deep_copy(execution_space(), devView, hostView);
    dv = dual_view_type(devView, hostView);
    execution_space().fence();
  }
}

template <class ElementType, class DeviceType>
void makeDualViewFromArrayView(Kokkos::DualView<ElementType*, DeviceType>& dv,
                               const Teuchos::ArrayView<const ElementType>& av,
                               const std::string& label) {
  using execution_space      = typename DeviceType::execution_space;
  using dual_view_type       = Kokkos::DualView<ElementType*, DeviceType>;
  using host_view_type       = typename dual_view_type::t_host;
  using const_host_view_type = typename host_view_type::const_type;

  const auto size        = av.size();
  const ElementType* ptr = (size == 0) ? nullptr : av.getRawPtr();
  const_host_view_type inView(ptr, size);
  host_view_type hostView(view_alloc_no_init(label), size);
  // DEEP_COPY REVIEW - DEVICE-TO-HOSTMIRROR
  Kokkos::deep_copy(execution_space(), hostView, inView);

  makeDualViewFromOwningHostView(dv, hostView);
}

template <class ElementType, class DeviceType>
void makeDualViewFromVector(Kokkos::DualView<ElementType*, DeviceType>& dv,
                            const std::vector<ElementType>& vec,
                            const std::string& label) {
  using dual_view_type       = Kokkos::DualView<ElementType*, DeviceType>;
  using execution_space      = typename DeviceType::execution_space;
  using host_view_type       = typename dual_view_type::t_host;
  using const_host_view_type = typename host_view_type::const_type;

  const auto size        = vec.size();
  const ElementType* ptr = (size == 0) ? nullptr : vec.data();
  const_host_view_type inView(ptr, size);
  host_view_type hostView(view_alloc_no_init(label), size);
  // DEEP_COPY REVIEW - DEVICE-TO-HOSTMIRROR
  Kokkos::deep_copy(execution_space(), hostView, inView);

  makeDualViewFromOwningHostView(dv, hostView);
}

template <class ElementType, class DeviceType>
void printDualView(std::ostream& out,
                   const Kokkos::DualView<ElementType*, DeviceType>& dv,
                   const std::string& name) {
  out << name << ": ";
  const size_t size   = size_t(dv.extent(0));
  const auto hostView = dv.view_host();

  out << "[";
  for (size_t k = 0; k < size; ++k) {
    out << hostView[k];
    if (k + size_t(1) < size) {
      out << ",";
    }
  }
  out << "]";
}

}  // namespace Details

}  // namespace Tpetra

#endif  // TPETRA_DETAILS_DUALVIEWUTIL_HPP