/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

/*
This testcase verifies the following scenarios
1. hipMemcpyAsync with kernel launch
2. H2D-D2D-D2H-H2PinnMem and device context change scenarios
3. This test launches multiple threads which uses same stream to deploy kernel
   and also launch hipMemcpyAsync() api. This test case is simulate the scenario
   reported in SWDEV-181598.
*/

#include <hip_test_common.hh>
#include <hip_test_kernels.hh>
#include <hip_test_checkers.hh>
#include <atomic>

static constexpr auto NUM_ELM{1024 * 1024};


static size_t N_ELMTS() {
  static const size_t val = isQuickLevel() ? (8 * 1024) : (32 * 1024);
  return val;
}
std::atomic<size_t> Thread_count{0};
static unsigned blocksPerCU{6};  // to hide latency
static unsigned threadsPerBlock{256};

template <typename T>
void Thread_func(T* A_d, T* B_d, T* C_d, T* C_h, size_t Nbytes, hipStream_t mystream) {
  unsigned blocks = HipTest::setNumBlocks(blocksPerCU, threadsPerBlock, N_ELMTS());
  hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks), dim3(threadsPerBlock), 0, mystream, A_d,
                     C_d, N_ELMTS());
  HIP_CHECK_THREAD(hipGetLastError());
  HIP_CHECK_THREAD(hipMemcpyAsync(C_h, C_d, Nbytes, hipMemcpyDeviceToHost, mystream));
  // The following two MemcpyAsync calls are for sole
  // purpose of loading stream with multiple async calls
  HIP_CHECK_THREAD(hipMemcpyAsync(B_d, A_d, Nbytes, hipMemcpyDeviceToDevice, mystream));
  HIP_CHECK_THREAD(hipMemcpyAsync(B_d, A_d, Nbytes, hipMemcpyDeviceToDevice, mystream));
  Thread_count++;
}

template <typename T> void Thread_func_MultiStream() {
  T *A_d{nullptr}, *B_d{nullptr}, *C_d{nullptr};
  T *A_h{nullptr}, *B_h{nullptr}, *C_h{nullptr};
  size_t Nbytes = N_ELMTS() * sizeof(T);
  unsigned blocks = HipTest::setNumBlocks(blocksPerCU, threadsPerBlock, N_ELMTS());

  HipTest::initArraysT(&A_d, &B_d, &C_d, &A_h, &B_h, &C_h, N_ELMTS(), false);
  hipStream_t mystream;
  HIP_CHECK_THREAD(hipStreamCreateWithFlags(&mystream, hipStreamNonBlocking));
  HIP_CHECK_THREAD(hipMemcpyAsync(A_d, A_h, Nbytes, hipMemcpyHostToDevice, mystream));
  hipLaunchKernelGGL((HipTest::vector_square), dim3(blocks), dim3(threadsPerBlock), 0, mystream,
                     A_d, C_d, N_ELMTS());
  HIP_CHECK_THREAD(hipGetLastError());
  HIP_CHECK_THREAD(hipMemcpyAsync(C_h, C_d, Nbytes, hipMemcpyDeviceToHost, mystream));
  // The following hipMemcpyAsync() is called only to
  // load stream with multiple Async calls
  HIP_CHECK_THREAD(hipMemcpyAsync(B_d, A_d, Nbytes, hipMemcpyDeviceToDevice, mystream));
  Thread_count++;

  HIP_CHECK_THREAD(hipStreamSynchronize(mystream));
  HIP_CHECK_THREAD(hipStreamDestroy(mystream));
  // Verifying result of the kernel computation
  for (size_t i = 0; i < N_ELMTS(); i++) {
    auto res = A_h[i] * A_h[i];
    REQUIRE_THREAD(res == C_h[i]);
  }

  // Releasing resources
  HipTest::freeArraysT<T>(A_d, B_d, C_d, A_h, B_h, C_h, false);
}

/*
This testcase verifies hipMemcpyAsync API
Initializes device variables
Launches kernel and performs the sum of device variables
copies the result to host variable and validates the result.
*/
/*
This testcase verifies the following scenarios
1. H2H,H2PinMem and PinnedMem2Host
2. H2D-D2D-D2H in same GPU
3. Pinned Host Memory to device variables in same GPU
4. Device context change
5. H2D-D2D-D2H peer GPU
*/
HIP_TEMPLATE_TEST_CASE(Unit_hipMemcpyAsync_H2H_H2D_D2H_H2PinMem, char, int, float, double) {
  TestType *A_d{nullptr}, *B_d{nullptr};
  TestType *A_h{nullptr}, *B_h{nullptr};
  TestType *A_Ph{nullptr}, *B_Ph{nullptr};
  HIP_CHECK(hipSetDevice(0));
  hipStream_t stream;
  HIP_CHECK(hipStreamCreate(&stream));
  HipTest::initArrays<TestType>(&A_d, &B_d, nullptr, &A_h, &B_h, nullptr,
                                NUM_ELM * sizeof(TestType));
  HipTest::initArrays<TestType>(nullptr, nullptr, nullptr, &A_Ph, &B_Ph, nullptr,
                                NUM_ELM * sizeof(TestType), true);

  SECTION("H2H, H2PinMem and PinMem2H") {
    HIP_CHECK(hipMemcpyAsync(B_h, A_h, NUM_ELM * sizeof(TestType), hipMemcpyHostToHost, stream));
    HIP_CHECK(hipMemcpyAsync(A_Ph, B_h, NUM_ELM * sizeof(TestType), hipMemcpyHostToHost, stream));
    HIP_CHECK(hipMemcpyAsync(B_Ph, A_Ph, NUM_ELM * sizeof(TestType), hipMemcpyHostToHost, stream));
    HIP_CHECK(hipStreamSynchronize(stream));
    HipTest::checkTest(A_h, B_Ph, NUM_ELM);
  }

  SECTION("H2D-D2D-D2H-SameGPU") {
    HIP_CHECK(hipMemcpyAsync(A_d, A_h, NUM_ELM * sizeof(TestType), hipMemcpyHostToDevice, stream));
    HIP_CHECK(
        hipMemcpyAsync(B_d, A_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToDevice, stream));
    HIP_CHECK(hipMemcpyAsync(B_h, B_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToHost, stream));
    HIP_CHECK(hipStreamSynchronize(stream));
    HipTest::checkTest(A_h, B_h, NUM_ELM);
  }

  SECTION("pH2D-D2D-D2pH-SameGPU") {
    HIP_CHECK(hipMemcpyAsync(A_d, A_Ph, NUM_ELM * sizeof(TestType), hipMemcpyHostToDevice, stream));
    HIP_CHECK(
        hipMemcpyAsync(B_d, A_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToDevice, stream));
    HIP_CHECK(hipMemcpyAsync(B_Ph, B_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToHost, stream));
    HIP_CHECK(hipStreamSynchronize(stream));
    HipTest::checkTest(A_Ph, B_Ph, NUM_ELM);
  }
  SECTION("H2D-D2D-D2H-DeviceContextChange") {
    int deviceCount = 0;
    HIP_CHECK(hipGetDeviceCount(&deviceCount));
    if (deviceCount < 2) {
      WARN("Skipping section: " << HipTest::SkipReason::kFewerThanTwoGpus);
    } else {
      int canAccessPeer = 0;
      HIP_CHECK(hipDeviceCanAccessPeer(&canAccessPeer, 0, 1));
      if (canAccessPeer) {
        HIP_CHECK(hipSetDevice(1));
        HIP_CHECK(
            hipMemcpyAsync(A_d, A_h, NUM_ELM * sizeof(TestType), hipMemcpyHostToDevice, stream));
        HIP_CHECK(
            hipMemcpyAsync(B_d, A_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToDevice, stream));
        HIP_CHECK(
            hipMemcpyAsync(B_h, B_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToHost, stream));
        HIP_CHECK(hipStreamSynchronize(stream));
        HipTest::checkTest(A_h, B_h, NUM_ELM);

      } else {
        WARN("Skipping section: " << HipTest::SkipReason::kPeerAccessUnavailable);
      }
    }
  }

  SECTION("H2D-D2D-D2H-PeerGPU") {
    int deviceCount = 0;
    HIP_CHECK(hipGetDeviceCount(&deviceCount));
    if (deviceCount < 2) {
      WARN("Skipping section: " << HipTest::SkipReason::kFewerThanTwoGpus);
    } else {
      int canAccessPeer = 0;
      HIP_CHECK(hipDeviceCanAccessPeer(&canAccessPeer, 0, 1));
      if (canAccessPeer) {
        HIP_CHECK(hipSetDevice(1));
        TestType* C_d{nullptr};
        HipTest::initArrays<TestType>(nullptr, nullptr, &C_d, nullptr, nullptr, nullptr,
                                      NUM_ELM * sizeof(TestType));
        HIP_CHECK(
            hipMemcpyAsync(A_d, A_h, NUM_ELM * sizeof(TestType), hipMemcpyHostToDevice, stream));
        HIP_CHECK(
            hipMemcpyAsync(C_d, A_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToDevice, stream));
        HIP_CHECK(
            hipMemcpyAsync(B_h, C_d, NUM_ELM * sizeof(TestType), hipMemcpyDeviceToHost, stream));
        HIP_CHECK(hipStreamSynchronize(stream));
        HipTest::checkTest(A_h, B_h, NUM_ELM);
        HIP_CHECK(hipFree(C_d));

      } else {
        WARN("Skipping section: " << HipTest::SkipReason::kPeerAccessUnavailable);
      }
    }
  }

  HIP_CHECK(hipStreamDestroy(stream));

  HipTest::freeArrays<TestType>(A_d, B_d, nullptr, A_h, B_h, nullptr, false);
  HipTest::freeArrays<TestType>(nullptr, nullptr, nullptr, A_Ph, B_Ph, nullptr, true);
}

// This test launches multiple threads which uses same stream to deploy kernel
// and also launch hipMemcpyAsync() api. This test case is simulate the scenario
// reported in SWDEV-181598

HIP_TEMPLATE_TEST_CASE(Unit_hipMemcpyAsync_hipMultiMemcpyMultiThread, int, float, double) {
  const int NUM_THREADS = isQuickLevel() ? 4 : 16;
  size_t Nbytes = N_ELMTS() * sizeof(TestType);

  int Data_mismatch = 0;
  hipStream_t mystream;
  TestType *A_d{nullptr}, *B_d{nullptr}, *C_d{nullptr};
  TestType *A_h{nullptr}, *B_h{nullptr}, *C_h{nullptr};

  HipTest::initArrays(&A_d, &B_d, &C_d, &A_h, &B_h, &C_h, N_ELMTS(), false);

  HIP_CHECK(hipStreamCreateWithFlags(&mystream, hipStreamNonBlocking));
  HIP_CHECK(hipMemcpyAsync(A_d, A_h, Nbytes, hipMemcpyHostToDevice, mystream));

  std::thread T[NUM_THREADS];
  for (int i = 0; i < NUM_THREADS; i++) {
    T[i] = std::thread(Thread_func<TestType>, A_d, B_d, C_d, C_h, Nbytes, mystream);
  }

  // Wait until all the threads finish their execution
  for (int i = 0; i < NUM_THREADS; i++) {
    T[i].join();
  }

  HIP_CHECK_THREAD_FINALIZE();

  HIP_CHECK(hipStreamSynchronize(mystream));
  HIP_CHECK(hipStreamDestroy(mystream));

  // Verifying the result of the kernel computation
  for (size_t i = 0; i < N_ELMTS(); i++) {
    if (C_h[i] != A_h[i] * A_h[i]) {
      Data_mismatch++;
    }
  }
  REQUIRE(Thread_count.load() == NUM_THREADS);
  REQUIRE(Data_mismatch == 0);
  HipTest::freeArrays<TestType>(A_d, B_d, C_d, A_h, B_h, C_h, false);
  Thread_count.exchange(0);
}

HIP_TEMPLATE_TEST_CASE(Unit_hipMemcpyAsync_hipMultiMemcpyMultiThreadMultiStream, int, float,
                   double) {
  const int NUM_THREADS = isQuickLevel() ? 4 : 16;
  std::thread T[NUM_THREADS];
  for (int i = 0; i < NUM_THREADS; i++) {
    T[i] = std::thread(Thread_func_MultiStream<TestType>);
  }

  // Wait until all the threads finish their execution
  for (int i = 0; i < NUM_THREADS; i++) {
    T[i].join();
  }

  HIP_CHECK_THREAD_FINALIZE();

  REQUIRE(Thread_count.load() == NUM_THREADS);
  Thread_count.exchange(0);
}

/*
This testcase verifies hipMemcpy API with pinnedMemory and hostRegister
along with kernel launches
*/

HIP_TEMPLATE_TEST_CASE(Unit_hipMemcpyAsync_PinnedRegMemWithKernelLaunch, int, float, double) {
  int numDevices = 0;
  HIP_CHECK(hipGetDeviceCount(&numDevices));
  if (numDevices < 2) {
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }
  {
    // 1 refers to pinned Memory
    // 2 refers to register Memory
    int MallocPinType = GENERATE(0, 1);
    size_t Nbytes = NUM_ELM * sizeof(TestType);
    unsigned blocks = HipTest::setNumBlocks(blocksPerCU, threadsPerBlock, NUM_ELM);

    TestType *A_d{nullptr}, *B_d{nullptr}, *C_d{nullptr};
    TestType *X_d{nullptr}, *Y_d{nullptr}, *Z_d{nullptr};
    TestType *A_h{nullptr}, *B_h{nullptr}, *C_h{nullptr};
    if (MallocPinType) {
      HipTest::initArrays(&A_d, &B_d, &C_d, &A_h, &B_h, &C_h, NUM_ELM, true);
    } else {
      A_h = reinterpret_cast<TestType*>(malloc(Nbytes));
      HIP_CHECK(hipHostRegister(A_h, Nbytes, hipHostRegisterDefault));
      B_h = reinterpret_cast<TestType*>(malloc(Nbytes));
      HIP_CHECK(hipHostRegister(B_h, Nbytes, hipHostRegisterDefault));
      C_h = reinterpret_cast<TestType*>(malloc(Nbytes));
      HIP_CHECK(hipHostRegister(C_h, Nbytes, hipHostRegisterDefault));
      HipTest::initArrays<TestType>(&A_d, &B_d, &C_d, nullptr, nullptr, nullptr, NUM_ELM, false);
      HipTest::setDefaultData<TestType>(NUM_ELM, A_h, B_h, C_h);
    }
    HIP_CHECK(hipMemcpy(A_d, A_h, Nbytes, hipMemcpyHostToDevice));
    HIP_CHECK(hipMemcpy(B_d, B_h, Nbytes, hipMemcpyHostToDevice));

    hipLaunchKernelGGL(HipTest::vectorADD, dim3(blocks), dim3(threadsPerBlock), 0, 0,
                       static_cast<const TestType*>(A_d), static_cast<const TestType*>(B_d), C_d,
                       NUM_ELM);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemcpy(C_h, C_d, Nbytes, hipMemcpyDeviceToHost));
    HipTest::checkVectorADD(A_h, B_h, C_h, NUM_ELM);

    unsigned int seed = time(0);
    HIP_CHECK(hipSetDevice(HipTest::RAND_R(&seed) % (numDevices - 1) + 1));

    int device;
    HIP_CHECK(hipGetDevice(&device));
    INFO("hipMemcpy is set to happen between device 0 and device " << device);
    HipTest::initArrays<TestType>(&X_d, &Y_d, &Z_d, nullptr, nullptr, nullptr, NUM_ELM, false);

    hipStream_t gpu1Stream;
    HIP_CHECK(hipStreamCreate(&gpu1Stream));

    for (int j = 0; j < NUM_ELM; j++) {
      A_h[j] = 0;
      B_h[j] = 0;
      C_h[j] = 0;
    }

    HIP_CHECK(hipMemcpy(A_h, A_d, Nbytes, hipMemcpyDeviceToHost));
    HIP_CHECK(hipMemcpyAsync(X_d, A_h, Nbytes, hipMemcpyHostToDevice, gpu1Stream));
    HIP_CHECK(hipMemcpy(B_h, B_d, Nbytes, hipMemcpyDeviceToHost));
    HIP_CHECK(hipMemcpyAsync(Y_d, B_h, Nbytes, hipMemcpyHostToDevice, gpu1Stream));

    hipLaunchKernelGGL(HipTest::vectorADD, dim3(blocks), dim3(threadsPerBlock), 0, 0,
                       static_cast<const TestType*>(X_d), static_cast<const TestType*>(Y_d), Z_d,
                       NUM_ELM);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemcpyAsync(C_h, Z_d, Nbytes, hipMemcpyDeviceToHost, gpu1Stream));
    HIP_CHECK(hipStreamSynchronize(gpu1Stream));

    HipTest::checkVectorADD(A_h, B_h, C_h, NUM_ELM);

    if (MallocPinType) {
      HipTest::freeArrays<TestType>(A_d, B_d, C_d, A_h, B_h, C_h, true);
    } else {
      HIP_CHECK(hipHostUnregister(A_h));
      free(A_h);
      HIP_CHECK(hipHostUnregister(B_h));
      free(B_h);
      HIP_CHECK(hipHostUnregister(C_h));
      free(C_h);
      HipTest::freeArrays<TestType>(A_d, B_d, C_d, nullptr, nullptr, nullptr, false);
    }
    HipTest::freeArrays<TestType>(X_d, Y_d, Z_d, nullptr, nullptr, nullptr, false);
    HIP_CHECK(hipStreamDestroy(gpu1Stream));
  }
}