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

#include <hip_test_kernels.hh>
#include <hip_test_checkers.hh>
#include <hip_test_common.hh>
#define GRIDSIZE 512
#define BLOCKSIZE 256
#define NUM_OF_STREAM 3
#define THREADS_PER_BLOCK 512
#define GRAPH_LAUNCH_ITERATIONS 1000

static __global__ void reduce(float* d_in, double* d_out) {
  int myId = threadIdx.x + blockDim.x * blockIdx.x;
  int tid = threadIdx.x;
  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
      d_in[myId] += d_in[myId + s];
    }
    __syncthreads();
  }
  if (tid == 0) {
    d_out[blockIdx.x] = d_in[myId];
  }
}
static __global__ void reduceFinal(double* d_in, double* d_out) {
  int myId = threadIdx.x + blockDim.x * blockIdx.x;
  int tid = threadIdx.x;
  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
      d_in[myId] += d_in[myId + s];
    }
    __syncthreads();
  }
  if (tid == 0) {
    *d_out = d_in[myId];
  }
}

static void init_input(float* a, size_t size) {
  unsigned int seed = time(nullptr);
  for (size_t i = 0; i < size; i++) {
    a[i] = (HipTest::RAND_R(&seed) & 0xFF) / static_cast<float>(RAND_MAX);
  }
}

static bool gPassed = true;
static void* gusrptr;
static void* ptr0xff = reinterpret_cast<void*>(0xffffffff);
static size_t NSize = GRIDSIZE * BLOCKSIZE;
static size_t Nbytes = NSize * sizeof(float);

typedef struct userDataStruct {
  float* A_h;
  float* C_h;
  float* A_d;
  float* C_d;
  bool isPassed;
  bool isOpCompleted;
} usrDataS;

// Common callback function.
static void Fn_validateSq(void* userData) {
  REQUIRE(userData != nullptr);
  usrDataS* ptrUsrData = reinterpret_cast<usrDataS*>(userData);
  for (size_t i = 0; i < NSize; i++) {
    if (ptrUsrData->C_h[i] != (ptrUsrData->A_h[i] * ptrUsrData->A_h[i])) {
      ptrUsrData->isPassed = false;
      return;
    }
  }
  ptrUsrData->isPassed = true;
}

// Test scenario 1
// simple scenario that validates passing userData to host function.
static void Fn_ChkUserdataPtr(void* userData) {
  gPassed = true;
  if (gusrptr != userData) {
    gPassed = false;
  }
}

HIP_TEST_CASE(Unit_hipLaunchHostFunc_basic) {
  hipStream_t mystream;
  HIP_CHECK(hipStreamCreate(&mystream));
  gusrptr = ptr0xff;
  gPassed = true;
  HIP_CHECK(hipLaunchHostFunc(mystream, Fn_ChkUserdataPtr, gusrptr));
  HIP_CHECK(hipStreamSynchronize(mystream));
  HIP_CHECK(hipStreamDestroy(mystream));
  REQUIRE(gPassed);
}

// Negative test scenario for hipLaunchHostFunc
HIP_TEST_CASE(Unit_hipLaunchHostFunc_Negative) {
  hipStream_t mystream;
  HIP_CHECK(hipStreamCreate(&mystream));

  SECTION("Pass nullptr as function") {
    REQUIRE(hipLaunchHostFunc(mystream, nullptr, 0) == hipErrorInvalidValue);
  }
  HIP_CHECK(hipStreamDestroy(mystream));
}

// Local Function
static void launchOperationOnStrm(usrDataS* usrDataptr, hipStream_t stream) {
  usrDataptr->isPassed = false;
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr->A_d)), Nbytes, stream));
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr->C_d)), Nbytes, stream));
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr->A_d, usrDataptr->A_h, Nbytes, hipMemcpyHostToDevice, stream));
  hipLaunchKernelGGL((HipTest::vector_square), dim3(GRIDSIZE), dim3(BLOCKSIZE), 0, stream,
                     usrDataptr->A_d, usrDataptr->C_d, NSize);
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr->C_h, usrDataptr->C_d, Nbytes, hipMemcpyDeviceToHost, stream));
  HIP_CHECK(hipLaunchHostFunc(stream, Fn_validateSq, reinterpret_cast<void*>(usrDataptr)));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr->A_d), stream));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr->C_d), stream));
  HIP_CHECK(hipStreamSynchronize(stream));
  REQUIRE(usrDataptr->isPassed);
}

// Test scenario 2
// scenario that validates the host launch function on 3 different streams,
// created stream, default/null stream and hipStreamPerThread.
HIP_TEST_CASE(Unit_hipLaunchHostFunc_streams) {
  hipStream_t stream[NUM_OF_STREAM];
  HIP_CHECK(hipStreamCreate(&stream[0]));
  stream[1] = 0;  // Null stream
  stream[2] = hipStreamPerThread;
  usrDataS* usrDataptr = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr != nullptr);
  usrDataptr->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->A_h != nullptr);
  usrDataptr->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->C_h != nullptr);
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr->A_h[i] = 21.0f;
  }
  for (int idx = 0; idx < NUM_OF_STREAM; idx++) {
    launchOperationOnStrm(usrDataptr, stream[idx]);
  }
  HIP_CHECK(hipStreamDestroy(stream[0]));
  free(usrDataptr->A_h);
  free(usrDataptr->C_h);
  free(usrDataptr);
}

// Test scenario 3
// test case to validate hipLaunchHostFunc for multi stream scenario.
// create 2 different streams and call hipLaunchHostFunc, stream synchronize.
static void Fn_validateMul_stream(void* userData) {
  REQUIRE(userData != nullptr);
  usrDataS* ptrUsrData = reinterpret_cast<usrDataS*>(userData);
  for (size_t i = 0; i < NSize; i++) {
    if (ptrUsrData->C_h[i] != (ptrUsrData->A_h[i] * ptrUsrData->A_h[i])) {
      ptrUsrData->isPassed = false;
      return;
    }
  }
  ptrUsrData->isPassed = true;
}

HIP_TEST_CASE(Unit_hipLaunchHostFunc_multistreams) {
  hipStream_t mystream1, mystream2;
  HIP_CHECK(hipStreamCreateWithFlags(&mystream1, hipStreamNonBlocking));
  HIP_CHECK(hipStreamCreateWithFlags(&mystream2, hipStreamNonBlocking));
  usrDataS* usrDataptr1 = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr1 != nullptr);
  usrDataS* usrDataptr2 = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr2 != nullptr);
  usrDataptr1->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr1->A_h != nullptr);
  usrDataptr1->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr1->C_h != nullptr);
  // input data
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr1->A_h[i] = 11.0f;
  }
  usrDataptr1->isPassed = false;
  usrDataptr2->isPassed = false;
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr1->A_d)), Nbytes, mystream1));
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr1->C_d)), Nbytes, mystream1));
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr1->A_d, usrDataptr1->A_h, Nbytes, hipMemcpyHostToDevice, mystream1));
  const unsigned blocks = GRIDSIZE;
  const unsigned threadsPerBlock = BLOCKSIZE;
  hipLaunchKernelGGL((HipTest::vector_square), dim3(blocks), dim3(threadsPerBlock), 0, mystream1,
                     usrDataptr1->A_d, usrDataptr1->C_d, NSize);
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr1->C_h, usrDataptr1->C_d, Nbytes, hipMemcpyDeviceToHost, mystream1));
  HIP_CHECK(
      hipLaunchHostFunc(mystream1, Fn_validateMul_stream, reinterpret_cast<void*>(usrDataptr1)));
  // launch kernel function for mystream2
  usrDataptr2->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr2->A_h != nullptr);
  usrDataptr2->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr2->C_h != nullptr);
  // input data
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr2->A_h[i] = 9.0f;
  }
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr2->A_d)), Nbytes, mystream2));
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr2->C_d)), Nbytes, mystream2));
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr2->A_d, usrDataptr2->A_h, Nbytes, hipMemcpyHostToDevice, mystream2));
  hipLaunchKernelGGL((HipTest::vector_square), dim3(blocks), dim3(threadsPerBlock), 0, mystream2,
                     usrDataptr2->A_d, usrDataptr2->C_d, NSize);
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr2->C_h, usrDataptr2->C_d, Nbytes, hipMemcpyDeviceToHost, mystream2));
  HIP_CHECK(
      hipLaunchHostFunc(mystream2, Fn_validateMul_stream, reinterpret_cast<void*>(usrDataptr2)));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr1->A_d), mystream1));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr1->C_d), mystream1));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr2->A_d), mystream2));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr2->C_d), mystream2));
  HIP_CHECK(hipStreamSynchronize(mystream1));
  HIP_CHECK(hipStreamSynchronize(mystream2));
  HIP_CHECK(hipStreamDestroy(mystream1));
  HIP_CHECK(hipStreamDestroy(mystream2));
  REQUIRE(usrDataptr1->isPassed);
  REQUIRE(usrDataptr2->isPassed);
  free(usrDataptr1->A_h);
  free(usrDataptr1->C_h);
  free(usrDataptr2->A_h);
  free(usrDataptr2->C_h);
  free(usrDataptr2);
  free(usrDataptr1);
}

// Test scenario 4
// test case to validate hipLaunchHostFunc for the kernel,
// validate hipLaunchHostFunc after kernel launch.
static void Fn_Completion_state(void* userData) {
  REQUIRE(userData != nullptr);
  usrDataS* ptrUsrData = reinterpret_cast<usrDataS*>(userData);
  ptrUsrData->isOpCompleted = true;
}

HIP_TEST_CASE(Unit_hipLaunchHostFunc_KernelHost) {
  hipStream_t stream1, stream2, stream3;
  HIP_CHECK(hipStreamCreate(&stream1));
  HIP_CHECK(hipStreamCreate(&stream2));
  HIP_CHECK(hipStreamCreate(&stream3));
  usrDataS* usrDataptr = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr != nullptr);
  usrDataptr->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->A_h != nullptr);
  usrDataptr->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->C_h != nullptr);
  // input data
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr->A_h[i] = 7.0f;
  }
  usrDataptr->isOpCompleted = false;
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr->A_d)), Nbytes, stream1));
  HIP_CHECK(hipMallocAsync(reinterpret_cast<void**>(&(usrDataptr->C_d)), Nbytes, stream1));
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr->A_d, usrDataptr->A_h, Nbytes, hipMemcpyHostToDevice, stream1));
  HIP_CHECK(hipLaunchHostFunc(stream1, Fn_Completion_state, reinterpret_cast<void*>(usrDataptr)));
  while (!usrDataptr->isOpCompleted) {
    std::this_thread::sleep_for(std::chrono::microseconds(100000));
  }  // Sleep for 100 ms*/
  usrDataptr->isOpCompleted = false;
  const unsigned blocks = GRIDSIZE;
  const unsigned threadsPerBlock = BLOCKSIZE;
  hipLaunchKernelGGL((HipTest::vector_square), dim3(blocks), dim3(threadsPerBlock), 0, stream2,
                     usrDataptr->A_d, usrDataptr->C_d, NSize);
  HIP_CHECK(hipLaunchHostFunc(stream2, Fn_Completion_state, reinterpret_cast<void*>(usrDataptr)));
  while (!usrDataptr->isOpCompleted) {
    std::this_thread::sleep_for(std::chrono::microseconds(100000));
  }  // Sleep for 100 ms*/
  usrDataptr->isOpCompleted = false;
  HIP_CHECK(
      hipMemcpyAsync(usrDataptr->C_h, usrDataptr->C_d, Nbytes, hipMemcpyDeviceToHost, stream3));
  HIP_CHECK(hipLaunchHostFunc(stream2, Fn_Completion_state, reinterpret_cast<void*>(usrDataptr)));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr->A_d), stream3));
  HIP_CHECK(hipFreeAsync(reinterpret_cast<void*>(usrDataptr->C_d), stream3));
  while (!usrDataptr->isOpCompleted) {
    std::this_thread::sleep_for(std::chrono::microseconds(100000));
  }  // Sleep for 100 ms*/
  for (size_t i = 0; i < NSize; i++) {
    if (usrDataptr->C_h[i] != (usrDataptr->A_h[i] * usrDataptr->A_h[i])) {
      REQUIRE(false);
    }
  }
  HIP_CHECK(hipStreamSynchronize(stream3));
  HIP_CHECK(hipStreamDestroy(stream3));
  HIP_CHECK(hipStreamDestroy(stream2));
  HIP_CHECK(hipStreamDestroy(stream1));
  free(usrDataptr->A_h);
  free(usrDataptr->C_h);
  free(usrDataptr);
}

// Test scenario 5
// scenario that validates the host launch function on multi device
// environment.
HIP_TEST_CASE(Unit_hipLaunchHostFunc_multidevice) {
  int num_devices;
  HIP_CHECK(hipGetDeviceCount(&num_devices));
  if (num_devices < 2) {
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }
  usrDataS* usrDataptr = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr != nullptr);
  usrDataptr->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->A_h != nullptr);
  usrDataptr->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->C_h != nullptr);
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr->A_h[i] = 21.0f;
  }
  for (int dev = 0; dev < num_devices; dev++) {
    HIP_CHECK(hipSetDevice(dev));
    hipStream_t stream;
    HIP_CHECK(hipStreamCreate(&stream));
    launchOperationOnStrm(usrDataptr, stream);
    HIP_CHECK(hipStreamDestroy(stream));
  }
  free(usrDataptr->A_h);
  free(usrDataptr->C_h);
  free(usrDataptr);
}

// Test scenario 6
// scenario that validates the host launch function on created
// stream with same priority.
HIP_TEST_CASE(Unit_hipLaunchHostFunc_Samepriority) {
  int priority = 0;
  unsigned int flags = 0;
  usrDataS* usrDataptr = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr != nullptr);
  usrDataptr->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->A_h != nullptr);
  usrDataptr->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->C_h != nullptr);
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr->A_h[i] = 21.0f;
  }
  for (int idx = 0; idx < NUM_OF_STREAM; idx++) {
    hipStream_t stream[NUM_OF_STREAM];
    HIP_CHECK(hipStreamCreateWithPriority(&stream[idx], flags, priority));
    launchOperationOnStrm(usrDataptr, stream[idx]);
    HIP_CHECK(hipStreamDestroy(stream[idx]));
  }
  free(usrDataptr->A_h);
  free(usrDataptr->C_h);
  free(usrDataptr);
}

// Test scenario 7
// scenario that validates the host launch function on
// created stream with different priority.
HIP_TEST_CASE(Unit_hipLaunchHostFunc_Diffpriority) {
  int priority;
  int priority_low{};
  int priority_high{};
  unsigned int flags = 0;
  HIP_CHECK(hipDeviceGetStreamPriorityRange(&priority_low, &priority_high));
  int numOfPriorities = priority_low - priority_high;
  const float arr_size = numOfPriorities + 1;
  hipStream_t* stream = reinterpret_cast<hipStream_t*>(malloc(arr_size * sizeof(hipStream_t)));
  stream[0] = 0;
  int count = 1;
  // Create a stream for each of the priority levels
  for (priority = priority_high; priority < priority_low; priority++) {
    HIP_CHECK(hipStreamCreateWithPriority(&stream[count++], flags, priority));
  }
  usrDataS* usrDataptr = reinterpret_cast<usrDataS*>(malloc(sizeof(usrDataS)));
  REQUIRE(usrDataptr != nullptr);
  usrDataptr->A_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->A_h != nullptr);
  usrDataptr->C_h = reinterpret_cast<float*>(malloc(Nbytes));
  REQUIRE(usrDataptr->C_h != nullptr);
  for (size_t i = 0; i < NSize; i++) {
    usrDataptr->A_h[i] = 11.0f;
  }
  for (int idx = 0; idx < arr_size; idx++) {
    launchOperationOnStrm(usrDataptr, stream[idx]);
  }
  count = 1;
  for (priority = priority_high; priority < priority_low; priority++) {
    HIP_CHECK(hipStreamDestroy(stream[count++]));
  }
  free(stream);
  free(usrDataptr->A_h);
  free(usrDataptr->C_h);
  free(usrDataptr);
}

// Test scenario 8
// create a graph by using hipGraphsUsingStreamCapture and call host function.

typedef struct callBackData {
  const char* fn_name;
  double* data;
} callBackData_t;
double result_gpu = 0.0;
void myHostNodeCallback(void* data) {
  static int iter = 0;
  iter++;
  // Check status of GPU after stream operations are done
  callBackData_t* tmp = reinterpret_cast<callBackData_t*>(data);
  // checkCudaErrors(tmp->status);
  double* result = reinterpret_cast<double*>(tmp->data);
  const char* function = reinterpret_cast<const char*>(tmp->fn_name);
  if (iter == GRAPH_LAUNCH_ITERATIONS)
    printf("[%s] Host callback final reduced sum = %lf\n", function, *result);
  result_gpu = *result;
  *result = 0.0;  // reset the result
}

HIP_TEST_CASE(Unit_hipLaunchHostFunc_Graph) {
  size_t size = 1 << 12;
  size_t maxBlocks = 512;
  float *inputVec_d = NULL, *inputVec_h = NULL;
  double *outputVec_d = NULL, *result_d;
  inputVec_h = reinterpret_cast<float*>(malloc(sizeof(float) * size));
  HIP_CHECK(hipMalloc(&inputVec_d, sizeof(float) * size));
  HIP_CHECK(hipMalloc(&outputVec_d, sizeof(double) * maxBlocks));
  HIP_CHECK(hipMalloc(&result_d, sizeof(double)));
  init_input(inputVec_h, size);
  hipStream_t stream1, stream2, stream3, streamForGraph;
  hipEvent_t forkStreamEvent, memsetEvent1, memsetEvent2;
  hipGraph_t graph;
  double result_h = 0.0;
  HIP_CHECK(hipStreamCreate(&stream1));
  HIP_CHECK(hipStreamCreate(&stream2));
  HIP_CHECK(hipStreamCreate(&stream3));
  HIP_CHECK(hipStreamCreate(&streamForGraph));
  HIP_CHECK(hipEventCreate(&forkStreamEvent));
  HIP_CHECK(hipEventCreate(&memsetEvent1));
  HIP_CHECK(hipEventCreate(&memsetEvent2));
  auto start = std::chrono::high_resolution_clock::now();
  HIP_CHECK(hipStreamBeginCapture(stream1, hipStreamCaptureModeGlobal));
  HIP_CHECK(hipEventRecord(forkStreamEvent, stream1));
  HIP_CHECK(hipStreamWaitEvent(stream2, forkStreamEvent, 0));
  HIP_CHECK(hipStreamWaitEvent(stream3, forkStreamEvent, 0));
  HIP_CHECK(
      hipMemcpyAsync(inputVec_d, inputVec_h, sizeof(float) * size, hipMemcpyDefault, stream1));
  HIP_CHECK(hipMemsetAsync(outputVec_d, 0, sizeof(double) * maxBlocks, stream2));
  HIP_CHECK(hipEventRecord(memsetEvent1, stream2));
  HIP_CHECK(hipMemsetAsync(result_d, 0, sizeof(double), stream3));
  HIP_CHECK(hipEventRecord(memsetEvent2, stream3));
  HIP_CHECK(hipStreamWaitEvent(stream1, memsetEvent1, 0));
  hipLaunchKernelGGL(reduce, dim3(size / THREADS_PER_BLOCK, 1, 1), dim3(THREADS_PER_BLOCK, 1, 1), 0,
                     stream1, inputVec_d, outputVec_d);
  HIP_CHECK(hipStreamWaitEvent(stream1, memsetEvent2, 0));
  hipLaunchKernelGGL(reduceFinal, dim3(1, 1, 1), dim3(THREADS_PER_BLOCK, 1, 1), 0, stream1,
                     outputVec_d, result_d);
  HIP_CHECK(hipMemcpyAsync(&result_h, result_d, sizeof(double), hipMemcpyDefault, stream1));

  callBackData_t hostFnData;
  hostFnData.data = &result_h;
  hostFnData.fn_name = "hipGraphsUsingStreamCapture";
  hipHostFn_t fn = myHostNodeCallback;
  HIP_CHECK(hipLaunchHostFunc(stream1, fn, &hostFnData));

  HIP_CHECK(hipStreamEndCapture(stream1, &graph));
  hipGraphNode_t* nodes = NULL;
  size_t numNodes = 0;
  HIP_CHECK(hipGraphGetNodes(graph, nodes, &numNodes));
  printf(
      "\nNum of nodes in the graph created using stream"
      "capture API = %zu\n",
      numNodes);
  HIP_CHECK(hipGraphGetRootNodes(graph, nodes, &numNodes));
  printf("root nodes in the graph created using stream capture API = %zu\n", numNodes);
  hipGraphExec_t graphExec;

  HIP_CHECK(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
  auto start1 = std::chrono::high_resolution_clock::now();
  const int graphIters = isQuickLevel() ? 10 : GRAPH_LAUNCH_ITERATIONS;
  for (int i = 0; i < graphIters; i++) {
    HIP_CHECK(hipGraphLaunch(graphExec, streamForGraph));
  }
  HIP_CHECK(hipStreamSynchronize(streamForGraph));
  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);
  std::cout << "Time taken for hipGraphsUsingStreamCapture with Init: "
            << std::chrono::duration_cast<std::chrono::milliseconds>(WithInit).count()
            << " milliseconds without Init:"
            << std::chrono::duration_cast<std::chrono::milliseconds>(WithoutInit).count()
            << " milliseconds " << std::endl;

  HIP_CHECK(hipGraphExecDestroy(graphExec));
  HIP_CHECK(hipGraphDestroy(graph));
  HIP_CHECK(hipEventDestroy(forkStreamEvent));
  HIP_CHECK(hipEventDestroy(memsetEvent1));
  HIP_CHECK(hipEventDestroy(memsetEvent2));
  HIP_CHECK(hipStreamDestroy(stream1));
  HIP_CHECK(hipStreamDestroy(stream2));
  HIP_CHECK(hipStreamDestroy(stream3));
  HIP_CHECK(hipStreamDestroy(streamForGraph));
  double result_h_cpu = 0.0;
  for (size_t i = 0; i < size; i++) {
    result_h_cpu += inputVec_h[i];
  }
  REQUIRE(result_h_cpu == result_gpu);
  HIP_CHECK(hipFree(inputVec_d));
  HIP_CHECK(hipFree(outputVec_d));
  HIP_CHECK(hipFree(result_d));
  free(inputVec_h);
}