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

#include <atomic>
#include <chrono>
#include <hip_test_common.hh>
#include <hip_test_defgroups.hh>
#include <unistd.h>
#include <vector>
#include "hipPerfCommon.hh"
/**
 * @addtogroup hipLaunchKernelGGL hipLaunchKernelGGL
 * @{
 * @ingroup PerformanceTestEvent
 */
__global__ void empty_kernel() {
  __shared__ int temp[256];
  temp[threadIdx.x] = sinf(float(threadIdx.x));
}
void rocm_empty_gpu_job(void* stream) {
  hipLaunchKernelGGL(empty_kernel, 1, 256, 0, (hipStream_t)stream);
}
std::vector<std::vector<hipStream_t>> stream_pools;
std::atomic<int> count(0);
std::atomic<bool> kill{false};
std::vector<PaddedTimestamp> thread_report;
void thread_jobs(int dev, int virt) {
  HIP_CHECK_PERF(hipSetDevice(dev));
  hipStream_t exec_stream = stream_pools[dev][virt];
  uint64_t n = 0;
  while (!kill) {
    rocm_empty_gpu_job(exec_stream);
    n++;
    if ((n & 150) == 0) {
      HIP_CHECK_PERF(hipStreamSynchronize(exec_stream));
      thread_report[dev * 4 + virt].ns.store(
          std::chrono::duration_cast<std::chrono::nanoseconds>(
              std::chrono::system_clock::now().time_since_epoch())
              .count(),
          std::memory_order_relaxed);
    }
    count++;
  }
  HIP_CHECK_PERF(hipStreamSynchronize(exec_stream));
}
/**
 * Test Description
 * ------------------------
 * - This test case prints the number of jobs/Second.
 * - 1) Launch number of threads on each device.
 * - 2) Threads should have only dispatches.
 * - 3) In the main thread calculate the number of jobs/Second
 * - 4) Print the jobs/Second value.
 * Test source
 * ------------------------
 * - performance/scenarios/event/hipKernelLookupPerf.cc
 * Test requirements
 * ------------------------
 * - HIP_VERSION >= 6.5
 */
HIP_TEST_CASE(Performance_hipKernelLookup) {
  int mgpu = 0;
  HIP_CHECK_PERF(hipGetDeviceCount(&mgpu));
  stream_pools.resize(mgpu);
  HIP_CHECK_PERF(hipSetDeviceFlags(hipDeviceScheduleSpin));
  for (int i = 0; i < mgpu; i++) {
    HIP_CHECK_PERF(hipSetDevice(i));
    stream_pools[i].resize(12);
    for (int j = 0; j < 12; j++)
      HIP_CHECK_PERF(hipStreamCreateWithFlags(&stream_pools[i][j], hipStreamNonBlocking));
  }
  // std::atomic is non-movable, so the vector cannot be resized — construct
  // it at the required size and move-assign it into the global.
  thread_report = std::vector<PaddedTimestamp>(mgpu * 4);
  for (int nDev = 1; nDev <= mgpu; nDev++) {
    count = 0;
    INFO("RUNNING ON " << nDev << " DEVICES\n");
    kill = false;
    std::vector<std::thread> threads;
    for (int i = 0; i < nDev * 4; i++) threads.push_back(std::thread(thread_jobs, i / 4, i % 4));
    usleep(1000000);
    auto t1 = std::chrono::system_clock::now();
    int counter = int(count);
    uint64_t total_count = 0;
    double total_time = 0;
    for (int t = 0; t < 10; t++) {
      usleep(1000000);
      auto t2 = std::chrono::system_clock::now();
      auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
      int counter2 = int(count);
      int64_t t2_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
                          t2.time_since_epoch())
                          .count();
      for (int i = 0; i < nDev * 4; i++) {
        int64_t slot_ns = thread_report[i].ns.load(std::memory_order_relaxed);
        if ((t2_ns - slot_ns) / 1000 >= 1000000) {
          INFO("Thread " << i / 4 << "/" << i % 4 << " is stuck\n");
        }
      }
      total_count += counter2 - counter;
      total_time += duration * 1e-6;
      t1 = t2;
      counter = counter2;
    }
    INFO("AVERAGE: " << total_count << "/" << total_time << " = " << total_count / total_time);
    kill = true;
    for (auto& t : threads) t.join();
    for (int i = 0; i < nDev; i++) {
      HIP_CHECK_PERF(hipSetDevice(i));
      HIP_CHECK_PERF(hipDeviceSynchronize());
    }
  }
  HIP_CHECK_PERF(hipSetDevice(0));
  for (int i = 0; i < mgpu; i++) {
    for (auto s : stream_pools[i]) HIP_CHECK_PERF(hipStreamDestroy(s));
  }
}
/**
 * End doxygen group PerformanceTest.
 * @}
 */