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

/*
 * Test for checking order of execution of device kernel and
 * hipMemsetAsync apis on all gpus
 */
#include <hip_test_common.hh>
#include <hip_test_checkers.hh>
#include <hip_test_kernels.hh>

#define ITER 6
#define N 1024 * 1024

constexpr auto blocksPerCU = 6;  // to hide latency
constexpr auto threadsPerBlock = 256;
static unsigned blocks = 0;


template <typename T> class MemSetKernelTest {
 public:
  T *A_h, *B_d, *B_h, *C_d;
  T memSetVal;
  size_t Nbytes;
  bool testResult = true;
  int validateCount = 0;
  hipStream_t stream;

  void memAllocate(T memSetValue) {
    memSetVal = memSetValue;
    Nbytes = N * sizeof(T);

    A_h = reinterpret_cast<T*>(malloc(Nbytes));
    HIP_ASSERT(A_h != nullptr);
    HIP_CHECK(hipMalloc(&B_d, Nbytes));
    B_h = reinterpret_cast<T*>(malloc(Nbytes));
    HIP_ASSERT(B_h != nullptr);
    HIP_CHECK(hipMalloc(&C_d, Nbytes));

    for (int i = 0; i < N; i++) {
      B_h[i] = i;
    }
    HIP_CHECK(hipMemcpy(B_d, B_h, Nbytes, hipMemcpyHostToDevice));
    HIP_CHECK(hipStreamCreate(&stream));
  }

  void memDeallocate() {
    HIP_CHECK(hipFree(B_d));
    HIP_CHECK(hipFree(C_d));
    free(B_h);
    free(A_h);
    HIP_CHECK(hipStreamDestroy(stream));
  }

  void validateExecutionOrder() {
    for (int p = 0; p < N; p++) {
      if (A_h[p] == memSetVal) {
        validateCount += 1;
      }
    }
  }

  bool resultAfterAllIterations() {
    testResult = (validateCount == (ITER * N)) ? true : false;
    memDeallocate();
    return testResult;
  }
};

static bool testhipMemsetAsyncWithKernel(bool UseStrmPerThrd) {
  MemSetKernelTest<char> obj;
  constexpr char memsetval = 0x42;

  obj.memAllocate(memsetval);
  for (int k = 0; k < ITER; ++k) {
    if (UseStrmPerThrd) {  // will use hipStreamPerThread stream object
      hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks), dim3(threadsPerBlock), 0,
                         hipStreamPerThread, obj.B_d, obj.C_d, N);
      HIP_CHECK(hipGetLastError());
      HIP_CHECK(hipStreamSynchronize(hipStreamPerThread));
      HIP_CHECK(hipMemsetAsync(obj.C_d, obj.memSetVal, N, hipStreamPerThread));
      HIP_CHECK(hipStreamSynchronize(hipStreamPerThread));
    } else {
      hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks), dim3(threadsPerBlock), 0, obj.stream,
                         obj.B_d, obj.C_d, N);
      HIP_CHECK(hipMemsetAsync(obj.C_d, obj.memSetVal, N, obj.stream));
      HIP_CHECK(hipStreamSynchronize(obj.stream));
    }
    HIP_CHECK(hipMemcpy(obj.A_h, obj.C_d, obj.Nbytes, hipMemcpyDeviceToHost));

    obj.validateExecutionOrder();
  }
  return obj.resultAfterAllIterations();
}

static bool testhipMemsetD32AsyncWithKernel() {
  MemSetKernelTest<int32_t> obj;
  constexpr int memsetD32val = 0xDEADBEEF;

  obj.memAllocate(memsetD32val);
  for (int k = 0; k < ITER; k++) {
    hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks), dim3(threadsPerBlock), 0, obj.stream,
                       obj.B_d, obj.C_d, N);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemsetD32Async((hipDeviceptr_t)obj.C_d, obj.memSetVal, N, obj.stream));
    HIP_CHECK(hipStreamSynchronize(obj.stream));
    HIP_CHECK(hipMemcpy(obj.A_h, obj.C_d, obj.Nbytes, hipMemcpyDeviceToHost));

    obj.validateExecutionOrder();
  }
  return obj.resultAfterAllIterations();
}

static bool testhipMemsetD16AsyncWithKernel() {
  MemSetKernelTest<int16_t> obj;
  constexpr int16_t memsetD16val = 0xDEAD;

  obj.memAllocate(memsetD16val);
  for (int k = 0; k < ITER; k++) {
    hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks), dim3(threadsPerBlock), 0, obj.stream,
                       obj.B_d, obj.C_d, N);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemsetD16Async((hipDeviceptr_t)obj.C_d, obj.memSetVal, N, obj.stream));
    HIP_CHECK(hipStreamSynchronize(obj.stream));
    HIP_CHECK(hipMemcpy(obj.A_h, obj.C_d, obj.Nbytes, hipMemcpyDeviceToHost));

    obj.validateExecutionOrder();
  }
  return obj.resultAfterAllIterations();
}

static bool testhipMemsetD8AsyncWithKernel() {
  MemSetKernelTest<char> obj;
  constexpr char memsetD8val = 0xDE;

  obj.memAllocate(memsetD8val);
  for (int k = 0; k < ITER; k++) {
    hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks), dim3(threadsPerBlock), 0, obj.stream,
                       obj.B_d, obj.C_d, N);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipMemsetD8Async((hipDeviceptr_t)obj.C_d, obj.memSetVal, N, obj.stream));
    HIP_CHECK(hipStreamSynchronize(obj.stream));
    HIP_CHECK(hipMemcpy(obj.A_h, obj.C_d, obj.Nbytes, hipMemcpyDeviceToHost));

    obj.validateExecutionOrder();
  }
  return obj.resultAfterAllIterations();
}


/*
 * Test for checking order of execution of device kernel and
 * hipMemsetAsync apis on all gpus
 */
HIP_TEST_CASE(Unit_hipMemsetAsync_VerifyExecutionWithKernel) {
  int numDevices = 0;
  bool ret, UseStrmPerThrd = false;

  blocks = HipTest::setNumBlocks(blocksPerCU, threadsPerBlock, N);

  HIP_CHECK(hipGetDeviceCount(&numDevices));
  REQUIRE(numDevices > 0);

  auto devNum = isQuickLevel() ? 0 : GENERATE_COPY(range(0, numDevices));
  HIP_CHECK(hipSetDevice(devNum));

  SECTION("hipMemsetAsync With Kernel") {
    UseStrmPerThrd = false;
    ret = testhipMemsetAsyncWithKernel(UseStrmPerThrd);
    REQUIRE(ret == true);
  }
  SECTION("hipMemsetAsync With Kernel using hipStreamPerThread stream obj") {
    UseStrmPerThrd = true;
    ret = testhipMemsetAsyncWithKernel(UseStrmPerThrd);
    REQUIRE(ret == true);
    HIP_CHECK(hipDeviceReset());
  }

  SECTION("hipMemsetD32Async With Kernel") {
    ret = testhipMemsetD32AsyncWithKernel();
    REQUIRE(ret == true);
  }

  SECTION("hipMemsetD16Async With Kernel") {
    ret = testhipMemsetD16AsyncWithKernel();
    REQUIRE(ret == true);
  }

  SECTION("hipMemsetD8Async With Kernel") {
    ret = testhipMemsetD8AsyncWithKernel();
    REQUIRE(ret == true);
  }
}