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

/**
 Simple test to demonstrate usage of graph.
 Compares implementation with and without using graphs.
*/

#include <hip_test_common.hh>

#ifdef KERNEL_ARG_PREFETCH
// Trimmed for mi4xx emu
#define N 32 * 32
#define NSTEP 1
#define NKERNEL 5
#else
#define N 1024 * 1024
#define NSTEP 1000
#define NKERNEL 25
#endif  // KERNEL_ARG_PREFETCH
#define CONSTANT 5.34

static __global__ void simpleKernel(float* out_d, float* in_d) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < N) out_d[idx] = CONSTANT * in_d[idx];
}

static void hipTestWithGraph() {
  int deviceId;
  HIP_CHECK(hipGetDevice(&deviceId));
  hipDeviceProp_t props;
  HIP_CHECK(hipGetDeviceProperties(&props, deviceId));

  hipStream_t stream;
  HIP_CHECK(hipStreamCreate(&stream));

  float *in_h, *out_h;
  in_h = new float[N];
  out_h = new float[N];
  for (int i = 0; i < N; i++) {
    in_h[i] = i;
  }

  float *in_d, *out_d;
  HIP_CHECK(hipMalloc(&in_d, N * sizeof(float)));
  HIP_CHECK(hipMalloc(&out_d, N * sizeof(float)));
  HIP_CHECK(hipMemcpy(in_d, in_h, N * sizeof(float), hipMemcpyHostToDevice));

  auto start = std::chrono::high_resolution_clock::now();
  // start CPU wallclock timer
  hipGraph_t graph;
  hipGraphExec_t instance;

  HIP_CHECK(hipStreamBeginCapture(stream, hipStreamCaptureModeGlobal));
  for (int ikrnl = 0; ikrnl < NKERNEL; ikrnl++) {
    simpleKernel<<<dim3(N / 512, 1, 1), dim3(512, 1, 1), 0, stream>>>(out_d, in_d);
  }
  HIP_CHECK(hipStreamEndCapture(stream, &graph));
  HIP_CHECK(hipGraphInstantiate(&instance, graph, nullptr, nullptr, 0));

  const int nstep = isQuickLevel() ? 10 : NSTEP;
  auto start1 = std::chrono::high_resolution_clock::now();
  for (int istep = 0; istep < nstep; istep++) {
    HIP_CHECK(hipGraphLaunch(instance, stream));
    HIP_CHECK(hipStreamSynchronize(stream));
  }
  auto stop = std::chrono::high_resolution_clock::now();
  auto withInit = std::chrono::duration<double, std::milli>(stop - start);
  auto withoutInit = std::chrono::duration<double, std::milli>(stop - start1);

  INFO("Time taken for graph with Init: "
       << std::chrono::duration_cast<std::chrono::milliseconds>(withInit).count()
       << " milliseconds without Init:"
       << std::chrono::duration_cast<std::chrono::milliseconds>(withoutInit).count()
       << " milliseconds ");

  HIP_CHECK(hipMemcpy(out_h, out_d, N * sizeof(float), hipMemcpyDeviceToHost));
  for (int i = 0; i < N; i++) {
    if (static_cast<float>(in_h[i] * CONSTANT) != out_h[i]) {
      INFO("Mismatch at indx:" << i << " " << in_h[i] << " " << out_h[i]);
      REQUIRE(false);
    }
  }
  delete[] in_h;
  delete[] out_h;
  HIP_CHECK(hipFree(in_d));
  HIP_CHECK(hipFree(out_d));
  HIP_CHECK(hipGraphExecDestroy(instance));
  HIP_CHECK(hipGraphDestroy(graph));
  HIP_CHECK(hipStreamDestroy(stream));
}

static void hipTestWithoutGraph() {
  int deviceId;
  HIP_CHECK(hipGetDevice(&deviceId));
  hipDeviceProp_t props;
  HIP_CHECK(hipGetDeviceProperties(&props, deviceId));
  INFO("Info: running on device " << deviceId << props.name);

  hipStream_t stream;
  HIP_CHECK(hipStreamCreate(&stream));

  float *in_h, *out_h;
  in_h = new float[N];
  out_h = new float[N];
  for (int i = 0; i < N; i++) {
    in_h[i] = i;
  }

  float *in_d, *out_d;
  HIP_CHECK(hipMalloc(&in_d, N * sizeof(float)));
  HIP_CHECK(hipMalloc(&out_d, N * sizeof(float)));
  HIP_CHECK(hipMemcpy(in_d, in_h, N * sizeof(float), hipMemcpyHostToDevice));

  // start CPU wallclock timer
  const int nstep2 = isQuickLevel() ? 10 : NSTEP;
  auto start = std::chrono::high_resolution_clock::now();
  for (int istep = 0; istep < nstep2; istep++) {
    for (int ikrnl = 0; ikrnl < NKERNEL; ikrnl++) {
      simpleKernel<<<dim3(N / 512, 1, 1), dim3(512, 1, 1), 0, stream>>>(out_d, in_d);
    }
    HIP_CHECK(hipStreamSynchronize(stream));
  }
  auto stop = std::chrono::high_resolution_clock::now();
  auto result = std::chrono::duration<double, std::milli>(stop - start);
  INFO("Time taken for test without graph: "
       << std::chrono::duration_cast<std::chrono::milliseconds>(result).count() << " millisecs ");
  HIP_CHECK(hipMemcpy(out_h, out_d, N * sizeof(float), hipMemcpyDeviceToHost));
  for (int i = 0; i < N; i++) {
    if (static_cast<float>(in_h[i] * CONSTANT) != out_h[i]) {
      INFO("Mismatch at indx:" << i << " " << in_h[i] << " " << out_h[i]);
      REQUIRE(false);
    }
  }
  delete[] in_h;
  delete[] out_h;
  HIP_CHECK(hipFree(in_d));
  HIP_CHECK(hipFree(out_d));
  HIP_CHECK(hipStreamDestroy(stream));
}

/**
 * Simple test to demonstrate usage of graph.
 */
#ifdef KERNEL_ARG_PREFETCH
TEST_CASE("Unit_hipGraph_SimpleGraphWithKernel_kernel_arg_prefetch") {
  if (!HipTest::isKernelArgPrefetchSupported()) {
    HIP_SKIP_TEST("Kernel arg prefetch is not supported on the device. Skipped.");
  }
#else
HIP_TEST_CASE(Unit_hipGraph_SimpleGraphWithKernel) {
#endif  // KERNEL_ARG_PREFETCH
  // Sections run test with and without graph.
  SECTION("Run Test Without Graph") { hipTestWithoutGraph(); }

  SECTION("Run Test With Graph") { hipTestWithGraph(); }
}