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


/* Test Description:
   Scenario-1: Launch a kernel in hipStreamPerThread, while it is in flight
   check for hipStreamQuery(hipStreamPerThread) it should return
   hipErrorNotReady.
   Scenario-2: Testing hipStreamPerThread stream object with hipMallocManaged()
   memory
   Scenario-3: To check the working of hipStreamPerThread in forked process
   Scenario-4: The following test case tests the working of hipEventSynchronize
   in multiple threads which are launched in quick succession
   Scenario-5: The following test case checks the working of
   hipStreamWaitEvent() with hipStreamWaitEvent()
   Scenario-6: Testing hipLaunchCooperativeKernel() api with hipStreamPerThread
   Scenario-7: Testing hipLaunchCooperativeKernelMultiDevice() with
   hipStreamPerThread
*/
#include <vector>
#include <thread>
#include <chrono>
#ifdef _WIN32
#include <Windows.h>
#define sleep(x) _sleep(x)
#endif
#ifdef __linux__
#include <unistd.h>
#include <sys/mman.h>
#include <sys/wait.h>
#endif

#include <hip_test_common.hh>
#include "hip/hip_cooperative_groups.h"
using namespace std::chrono;
using namespace cooperative_groups;
#if HT_AMD
#define HIPRT_CB
#endif


static bool IfTestPassed = false;
// kernel
__global__ void StreamPerThrd(int* Ad, int* Ad1, size_t n, int Pk_Clk, int Wait, int WaitEvnt = 0) {
  size_t index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < n) {
    Ad[index] = Ad[index] + 10;
  }
  if (Wait) {
    int64_t GpuFrq = (Pk_Clk * 1000);
    int64_t StrtTck = clock64();
    if (index == 0) {
      // The following while loop holds the execution for ~2 seconds.
      // Busy sleep on nvidia
      while ((clock64() - StrtTck) <= (2 * GpuFrq)) {
#if HT_AMD
        __builtin_amdgcn_s_sleep(10);
#endif
      }
      if (WaitEvnt == 1) {
        *Ad1 = 1;
      }
    }
  }
}


__global__ void StreamPerThrd1(int* A, int Pk_Clk) {
  int64_t GpuFrq = (Pk_Clk * 1000);
  int64_t StrtTck = clock64();
  // The following while loop holds the execution for ~1 second
  // Busy sleep on nvidia
  while ((clock64() - StrtTck) <= (GpuFrq)) {
#if HT_AMD
    __builtin_amdgcn_s_sleep(10);
#endif
  }
  *A = 1;
}

__global__ void StreamPerThrd_gfx11(int* Ad, int* Ad1, size_t n, int Pk_Clk, int Wait,
                                    int WaitEvnt = 0) {
#if HT_AMD
  size_t index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < n) {
    Ad[index] = Ad[index] + 10;
  }
  if (Wait) {
    int64_t GpuFrq = (Pk_Clk * 1000);
    int64_t StrtTck = clock_function();
    if (index == 0) {
      // The following while loop checks the value in ptr for around 4 seconds
      while ((clock_function() - StrtTck) <= (6 * GpuFrq)) {
      }
      if (WaitEvnt == 1) {
        *Ad1 = 1;
      }
    }
  }
#endif
}

__global__ void StreamPerThrd1_gfx11(int* A, int Pk_Clk) {
#if HT_AMD
  int64_t GpuFrq = (Pk_Clk * 1000);
  int64_t StrtTck = clock_function();
  // The following while loop checks the value in ptr for around 3-4 seconds
  while ((clock_function() - StrtTck) <= (3 * GpuFrq)) {
  }
  *A = 1;
#endif
}

__global__ void MiniKernel(int* A) {
  if (*A == 0) {
    *A = 2;  //  Fail condition
  } else if (*A == 1) {
    *A = 3;  //  Pass condition
  } else {
    *A = 4;  //  Garbage value found in A
  }
}

__global__ void StreamPerThrdCoopKrnl(int* Ad, int* n) {
  int NumElms = (*n);
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < NumElms) {
    Ad[index] = Ad[index] + 10;
  }
}

#if HT_AMD
__global__ void test_gwsPerThrd(uint* buf, uint bufSize, int64_t* tmpBuf, int64_t* result) {
  extern __shared__ int64_t tmp[];
  uint groups = gridDim.x;
  uint group_id = blockIdx.x;
  uint local_id = threadIdx.x;
  uint chunk = gridDim.x * blockDim.x;

  uint i = group_id * blockDim.x + local_id;
  int64_t sum = 0;
  while (i < bufSize) {
    sum += buf[i];
    i += chunk;
  }
  tmp[local_id] = sum;
  __syncthreads();
  i = 0;
  if (local_id == 0) {
    sum = 0;
    while (i < blockDim.x) {
      sum += tmp[i];
      i++;
    }
    tmpBuf[group_id] = sum;
  }

  // wait
  cooperative_groups::this_grid().sync();

  if (((blockIdx.x * blockDim.x) + threadIdx.x) == 0) {
    for (uint i = 1; i < groups; ++i) {
      sum += tmpBuf[i];
    }
    // *result = sum;
    result[1 + cooperative_groups::this_multi_grid().grid_rank()] = sum;
  }
  cooperative_groups::this_multi_grid().sync();
  if (cooperative_groups::this_multi_grid().grid_rank() == 0) {
    sum = 0;
    for (uint i = 1; i <= cooperative_groups::this_multi_grid().num_grids(); ++i) {
      sum += result[i];
    }
    *result = sum;
  }
}
#endif

// callback function
static void HIPRT_CB CallBackFunctn(hipStream_t strm, hipError_t err, void* ChkVal) {
  // The following HIPASSERT() is just to satisfy catch2 framework.
  // As it ensures the use of all the variables.
  HIPASSERT(strm);
  HIPCHECK(err);
  if (*(reinterpret_cast<int*>(ChkVal)) == 1) {
    IfTestPassed = true;
  } else {
    IfTestPassed = false;
  }
}

static void EventSync() {
  int *Ad = nullptr, *Ah = nullptr, NumElms = 4096, CONST_NUM = 123;
  int blockSize = 32, peak_clk;
  HIP_CHECK(hipMalloc(&Ad, NumElms * sizeof(int)));
  Ah = new int[NumElms];
  for (int i = 0; i < NumElms; ++i) {
    Ah[i] = CONST_NUM;
  }
  // creating event objects
  hipEvent_t start, end;
  HIP_CHECK(hipEventCreate(&start));
  HIP_CHECK(hipEventCreate(&end));
  HIP_CHECK(hipMemcpy(Ad, Ah, NumElms * sizeof(int), hipMemcpyHostToDevice));
  dim3 dimBlock(blockSize, 1, 1);
  dim3 dimGrid((NumElms + blockSize - 1) / blockSize, 1, 1);
  HIP_CHECK(hipEventRecord(start, hipStreamPerThread));
  if (IsGfx11()) {
    HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeWallClockRate, 0));
    StreamPerThrd_gfx11<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Ad, NULL, NumElms, peak_clk,
                                                                      0);
  } else {
    HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeClockRate, 0));
    StreamPerThrd<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Ad, NULL, NumElms, peak_clk, 0);
  }
  HIP_CHECK(hipEventRecord(end, hipStreamPerThread));
  HIP_CHECK(hipEventSynchronize(end));
  HIP_CHECK(hipMemcpy(Ah, Ad, NumElms * sizeof(int), hipMemcpyDeviceToHost));
  int MisMatch = 0;
  for (int i = 0; i < NumElms; ++i) {
    if (Ah[i] != (CONST_NUM + 10)) {
      MisMatch++;
    }
  }
  delete[] Ah;
  HIP_CHECK(hipFree(Ad));
  if (MisMatch) {
    WARN("Data Mismatch observed!!\n");
    IfTestPassed = false;
  } else {
    IfTestPassed = true;
  }

  HIP_CHECK(hipEventDestroy(start));
  HIP_CHECK(hipEventDestroy(end));
}

/* Launch a kernel in hipStreamPerThread, while it is in flight check for
   hipStreamQuery(hipStreamPerThread) it should return hipErrorNotReady.*/
HIP_TEST_CASE(Unit_hipStreamPerThreadTst_StrmQuery) {
  int *Ad = nullptr, *Ah = nullptr, NumElms = 4096, CONST_NUM = 123;
  int blockSize = 32, peak_clk;
  hipError_t err;
  HIP_CHECK(hipMalloc(&Ad, NumElms * sizeof(int)));
  Ah = new int[NumElms];
  for (int i = 0; i < NumElms; ++i) {
    Ah[i] = CONST_NUM;
  }
  HIP_CHECK(hipMemcpy(Ad, Ah, NumElms * sizeof(int), hipMemcpyHostToDevice));
  dim3 dimBlock(blockSize, 1, 1);
  dim3 dimGrid((NumElms + blockSize - 1) / blockSize, 1, 1);
  SECTION("Test working of hipStreamQuery") {
    if (IsGfx11()) {
      HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeWallClockRate, 0));
      StreamPerThrd_gfx11<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Ad, NULL, NumElms, peak_clk,
                                                                        1);
    } else {
      HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeClockRate, 0));
      StreamPerThrd<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Ad, NULL, NumElms, peak_clk, 1);
    }
    err = hipStreamQuery(hipStreamPerThread);
    if (err != hipErrorNotReady) {
      WARN("hipStreamQuery on hipStreamPerThread didnt return expected error!");
      IfTestPassed = false;
    } else {
      IfTestPassed = true;
    }
  }
  SECTION("check working of hipStreamAddCallback() with hipStreamPerThread") {
    int *Hptr = nullptr, *A_d = nullptr;
    HIP_CHECK(hipHostMalloc(&Hptr, sizeof(int)));
    *Hptr = 0;
    HIP_CHECK(hipHostGetDevicePointer(reinterpret_cast<void**>(&A_d), Hptr, 0));
    if (IsGfx11()) {
      StreamPerThrd1_gfx11<<<1, 1, 0, hipStreamPerThread>>>(A_d, peak_clk);
    } else {
      StreamPerThrd1<<<1, 1, 0, hipStreamPerThread>>>(A_d, peak_clk);
    }
    HIP_CHECK(hipStreamAddCallback(hipStreamPerThread, CallBackFunctn, A_d, 0));
    HIP_CHECK(hipStreamSynchronize(hipStreamPerThread));
    HIP_CHECK(hipHostFree(Hptr));
  }
  HIP_CHECK(hipFree(Ad));
  delete[] Ah;
  REQUIRE(IfTestPassed);
}

/* Testing hipStreamPerThread stream object with hipMallocManaged() memory*/
HIP_TEST_CASE(Unit_hipStreamPerThread_MangdMem) {
  CHECK_MANAGED_MEMORY_SUPPORT

  int *Hmm = nullptr, NumElms = 4096, CONST_NUM = 123, blockSize = 32;
  SECTION("Using Managed memory") {
    HIP_CHECK(hipMallocManaged(&Hmm, NumElms * sizeof(int)));
    for (int i = 0; i < NumElms; ++i) {
      Hmm[i] = CONST_NUM;
    }
  }
  SECTION("Prefetching Managed memory to device") {
    HIP_CHECK(hipMallocManaged(&Hmm, NumElms * sizeof(int)));
    for (int i = 0; i < NumElms; ++i) {
      Hmm[i] = CONST_NUM;
    }
    HIP_CHECK(hipMemPrefetchAsync(Hmm, NumElms * sizeof(int), 0, hipStreamPerThread));
  }
  int peak_clk;
  dim3 dimBlock(blockSize, 1, 1);
  dim3 dimGrid((NumElms + blockSize - 1) / blockSize, 1, 1);
  if (IsGfx11()) {
    HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeWallClockRate, 0));
    StreamPerThrd_gfx11<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Hmm, NULL, NumElms, peak_clk,
                                                                      0);
  } else {
    HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeClockRate, 0));
    StreamPerThrd<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Hmm, NULL, NumElms, peak_clk, 0);
  }
  HIP_CHECK(hipStreamSynchronize(hipStreamPerThread));
  // Validating the result
  int MisMatch = 0;
  for (int i = 0; i < NumElms; ++i) {
    if (Hmm[i] != (CONST_NUM + 10)) {
      MisMatch++;
    }
  }
  HIP_CHECK(hipFree(Hmm));
  REQUIRE(MisMatch == 0);
}

/*  To check the working of hipStreamPerThread in forked process*/
#ifdef __linux__
HIP_TEST_CASE(Unit_hipStreamPerThread_ChildProc) {
  if (fork() == 0) {  //  child process
    int *Ad = nullptr, *Ah = nullptr, NumElms = 4096, CONST_NUM = 123;
    int blockSize = 32, peak_clk;
    HIP_CHECK(hipMalloc(&Ad, NumElms * sizeof(int)));
    Ah = new int[NumElms];
    for (int i = 0; i < NumElms; ++i) {
      Ah[i] = CONST_NUM;
    }
    HIP_CHECK(hipMemcpy(Ad, Ah, NumElms * sizeof(int), hipMemcpyHostToDevice));
    dim3 dimBlock(blockSize, 1, 1);
    dim3 dimGrid((NumElms + blockSize - 1) / blockSize, 1, 1);
    if (IsGfx11()) {
      HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeWallClockRate, 0));
      StreamPerThrd_gfx11<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Ad, NULL, NumElms, peak_clk,
                                                                        0);
    } else {
      HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeClockRate, 0));
      StreamPerThrd<<<dimGrid, dimBlock, 0, hipStreamPerThread>>>(Ad, NULL, NumElms, peak_clk, 0);
    }
    HIP_CHECK(hipStreamSynchronize(hipStreamPerThread));
    HIP_CHECK(hipMemcpy(Ah, Ad, NumElms * sizeof(int), hipMemcpyDeviceToHost));
    int MisMatch = 0;
    for (int i = 0; i < NumElms; ++i) {
      if (Ah[i] != (CONST_NUM + 10)) {
        MisMatch++;
      }
    }
    delete[] Ah;
    HIP_CHECK(hipFree(Ad));
    if (MisMatch) {
      WARN("Data Mismatch observed!!\n");
      exit(9);
    } else {
      exit(10);
    }
  } else {  //  Parent process
    int stat;
    wait(&stat);
    int Result = WEXITSTATUS(stat);
    if (Result != 10) {
      REQUIRE(false);
    }
  }
}
#endif

/* The following test case tests the working of hipEventSynchronize in
   multiple threads which are launched in quick succession*/
HIP_TEST_CASE(Unit_hipStreamPerThread_EvtRcrdMThrd) {
  IfTestPassed = true;
  int MAX_THREAD_CNT = 20;
  std::vector<std::thread> threads(MAX_THREAD_CNT);
  for (auto& th : threads) {
    th = std::thread(EventSync);
  }
  for (auto& th : threads) {
    th.join();
  }
  REQUIRE(IfTestPassed);
}

/* The following test case checks the working of hipStreamWaitEvent() with
   hipStreamWaitEvent()*/
HIP_TEST_CASE(Unit_hipStreamPerThread_StrmWaitEvt) {
  IfTestPassed = true;
  int *Ad = nullptr, NumElms = 4096, CONST_NUM = 123, blockSize = 32, *Ah = nullptr;
  int *Ad1 = nullptr, *Ah1 = nullptr;
  Ah = new int[NumElms];
  Ah1 = new int;
  hipStream_t Strm;
  HIP_CHECK(hipStreamCreate(&Strm));
  for (int i = 0; i < NumElms; ++i) {
    Ah[i] = CONST_NUM;
  }
  Ah1[0] = 0;
  HIP_CHECK(hipMalloc(&Ad, NumElms * sizeof(int)));
  HIP_CHECK(hipMemcpy(Ad, Ah, NumElms * sizeof(int), hipMemcpyHostToDevice));
  memset(Ah, 0, NumElms * sizeof(int));
  HIP_CHECK(hipMalloc(&Ad1, sizeof(int)));
  HIP_CHECK(hipMemset(Ad1, 0, sizeof(int)));
  int peak_clk;
  dim3 dimBlock(blockSize, 1, 1);
  dim3 dimGrid((NumElms + blockSize - 1) / blockSize, 1, 1);
  hipEvent_t e1;
  HIPCHECK(hipEventCreate(&e1));
  if (IsGfx11()) {
    HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeWallClockRate, 0));
    StreamPerThrd_gfx11<<<dimGrid, dimBlock, 0, Strm>>>(Ad, Ad1, NumElms, peak_clk, 1, 1);
  } else {
    HIP_CHECK(hipDeviceGetAttribute(&peak_clk, hipDeviceAttributeClockRate, 0));
    StreamPerThrd<<<dimGrid, dimBlock, 0, Strm>>>(Ad, Ad1, NumElms, peak_clk, 1, 1);
  }
  HIP_CHECK(hipEventRecord(e1, Strm));
  HIP_CHECK(hipStreamWaitEvent(hipStreamPerThread, e1, 0 /*flags*/));
  MiniKernel<<<1, 1, 0, hipStreamPerThread>>>(Ad1);
  sleep(1);
  HIP_CHECK(hipMemcpy(Ah1, Ad1, sizeof(int), hipMemcpyDeviceToHost));
  if (*Ah1 != 3) {
    IfTestPassed = false;
    if (*Ah1 == 2) {
      WARN("hipStreamPerThread didn't honour hipStreamWaitEvent()");
    } else if (*Ah1 == 4) {
      WARN("Unexpected behavior observed with hipStreamPerThread");
    }
  }
  // Validating the result
  HIP_CHECK(hipMemcpy(Ah, Ad, NumElms * sizeof(int), hipMemcpyDeviceToHost));
  int MisMatch = 0;
  for (int i = 0; i < NumElms; ++i) {
    if (Ah[i] != (CONST_NUM + 10)) {
      MisMatch++;
    }
  }
  HIP_CHECK(hipFree(Ad));
  HIP_CHECK(hipFree(Ad1));
  HIP_CHECK(hipEventDestroy(e1));
  HIP_CHECK(hipStreamDestroy(Strm));
  delete[] Ah;
  delete Ah1;
  if (MisMatch) {
    WARN("Data mismatch observed!!\n");
    IfTestPassed = false;
  }
  REQUIRE(IfTestPassed);
}


/* Testing hipLaunchCooperativeKernel() api with hipStreamPerThread*/
HIP_TEST_CASE(Unit_hipStreamPerThread_CoopLaunch) {
  hipDeviceProp_t device_properties;
  HIPCHECK(hipGetDeviceProperties(&device_properties, 0));
  /* Test whether target device supports cooperative groups ****************/
  if (device_properties.cooperativeLaunch == 0) {
    HIP_SKIP_TEST(HipTest::SkipReason::kCooperativeLaunchUnsupported);
  } else {
    /* We will launch enough waves to fill up all of the GPU *****************/
    int warp_size = device_properties.warpSize;
    int num_sms = device_properties.multiProcessorCount;
    // long long totalTicks = device_properties.clockRate ;
    int max_blocks_per_sm = 0;
    // Calculate the device occupancy to know how many blocks can be run.
    HIPCHECK(hipOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, StreamPerThrdCoopKrnl,
                                                          warp_size, 0));
    int max_active_blocks = max_blocks_per_sm * num_sms;
    int *Ad = nullptr, *Ah = nullptr, *DNumElms = nullptr, NumElms = 4096;
    NumElms = max_active_blocks * warp_size > NumElms ? NumElms : max_active_blocks * warp_size;
    int Const = 123;
    Ah = new int[NumElms];
    for (int i = 0; i < NumElms; ++i) {
      Ah[i] = Const;
    }
    HIP_CHECK(hipMalloc(&Ad, sizeof(int) * NumElms));
    HIP_CHECK(hipMalloc(&DNumElms, sizeof(int)));
    HIP_CHECK(
        hipMemcpyAsync(Ad, Ah, sizeof(int) * NumElms, hipMemcpyHostToDevice, hipStreamPerThread));
    HIP_CHECK(
        hipMemcpyAsync(DNumElms, &NumElms, sizeof(int), hipMemcpyHostToDevice, hipStreamPerThread));
    HIP_CHECK(hipStreamSynchronize(hipStreamPerThread));

    void* coop_params[2];
    coop_params[0] = reinterpret_cast<void*>(&Ad);
    coop_params[1] = reinterpret_cast<void*>(&DNumElms);
    HIP_CHECK(hipLaunchCooperativeKernel(reinterpret_cast<void*>(StreamPerThrdCoopKrnl),
                                         max_active_blocks, warp_size, coop_params, 0,
                                         hipStreamPerThread));
    HIP_CHECK(hipMemcpy(Ah, Ad, sizeof(int) * NumElms, hipMemcpyDeviceToHost));
    // Verifying the result
    int DataMismatch = 0;
    for (int i = 0; i < NumElms; ++i) {
      if (Ah[i] != (Const + 10)) {
        DataMismatch++;
      }
    }
    HIP_CHECK(hipFree(Ad));
    HIP_CHECK(hipFree(DNumElms));
    delete[] Ah;
    if (DataMismatch > 0) {
      REQUIRE(false);
    }
  }
}