/*
//@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
*/

#include <Kokkos_Core.hpp>
#include <cuda/TestCuda_Category.hpp>

namespace Test {

template <class View>
struct CopyFunctor {
  View a;
  View b;

  CopyFunctor(int N) : a(View("A", N)), b(View("B", N)) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(int i) const { a(i) = b(i); }

  double time_copy(int R) {
    Kokkos::parallel_for("CopyFunctor::time_copy", a.extent(0), *this);
    Kokkos::fence();

    Kokkos::Timer timer;
    for (int r = 0; r < R; r++)
      Kokkos::parallel_for("CopyFunctor::time_copy", a.extent(0), *this);
    Kokkos::fence();
    return timer.seconds();
  }
};

TEST(cuda, debug_pin_um_to_host) {
  double time_cuda_space;
  double time_cuda_host_pinned_space;
  double time_cuda_uvm_space_not_pinned_1;
  double time_cuda_uvm_space_pinned;
  double time_cuda_uvm_space_not_pinned_2;

  int N = 10000000;
  int R = 100;
  {
    CopyFunctor<Kokkos::View<int*, Kokkos::CudaSpace>> f(N);
    time_cuda_space = f.time_copy(R);
  }
  {
    CopyFunctor<Kokkos::View<int*, Kokkos::CudaHostPinnedSpace>> f(N);
    time_cuda_host_pinned_space = f.time_copy(R);
  }
  {
    CopyFunctor<Kokkos::View<int*, Kokkos::CudaUVMSpace>> f(N);
    time_cuda_uvm_space_not_pinned_1 = f.time_copy(R);
  }
  {
#ifdef KOKKOS_IMPL_DEBUG_CUDA_PIN_UVM_TO_HOST
    kokkos_impl_cuda_set_pin_uvm_to_host(true);
#endif
    CopyFunctor<Kokkos::View<int*, Kokkos::CudaUVMSpace>> f(N);
    time_cuda_uvm_space_pinned = f.time_copy(R);
#ifdef KOKKOS_IMPL_DEBUG_CUDA_PIN_UVM_TO_HOST
    kokkos_impl_cuda_set_pin_uvm_to_host(false);
#endif
  }
  {
    CopyFunctor<Kokkos::View<int*, Kokkos::CudaUVMSpace>> f(N);
    time_cuda_uvm_space_not_pinned_2 = f.time_copy(R);
  }
  bool uvm_approx_cuda_1 =
      time_cuda_uvm_space_not_pinned_1 < time_cuda_space * 2.0;
  bool uvm_approx_cuda_2 =
      time_cuda_uvm_space_not_pinned_2 < time_cuda_space * 2.0;
  bool pinned_slower_cuda = time_cuda_host_pinned_space > time_cuda_space * 2.0;
  bool uvm_pinned_slower_cuda =
      time_cuda_uvm_space_pinned > time_cuda_space * 2.0;

  bool passed = uvm_approx_cuda_1 && uvm_approx_cuda_2 && pinned_slower_cuda &&
#ifdef KOKKOS_IMPL_DEBUG_CUDA_PIN_UVM_TO_HOST
                uvm_pinned_slower_cuda;
#else
                !uvm_pinned_slower_cuda;
#endif
  if (!passed)
    printf(
        "Time CudaSpace: %lf CudaUVMSpace_1: %lf CudaUVMSpace_2: %lf "
        "CudaPinnedHostSpace: %lf CudaUVMSpace_Pinned: %lf\n",
        time_cuda_space, time_cuda_uvm_space_not_pinned_1,
        time_cuda_uvm_space_not_pinned_2, time_cuda_host_pinned_space,
        time_cuda_uvm_space_pinned);
  ASSERT_TRUE(passed);
}

}  // namespace Test
