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

/* Test Case Description:
   1) This testcase verifies the hipMallocManaged basic scenario - supported on
     all devices
   2) This testcase verifies the hipMallocManaged advanced scenario - supported
     only on HMM enabled devices
   3) This testcase verifies that hipMallocManaged returns an OutOfMemory error
     for allocations much larger than the available memory - supported on all devices
*/

#include <hip_test_common.hh>
#include <hip_test_kernels.hh>
#include <hip_test_checkers.hh>
#include <string>


// Kernel functions

__global__ void KernelMul_MngdMem(int* Hmm, int* Dptr, size_t n) {
  size_t index = blockIdx.x * blockDim.x + threadIdx.x;
  size_t stride = blockDim.x * gridDim.x;
  for (size_t i = index; i < n; i += stride) {
    Hmm[i] = Dptr[i] * 10;
  }
}

__global__ void KernelMulAdd_MngdMem(int* Hmm, size_t n) {
  size_t index = blockIdx.x * blockDim.x + threadIdx.x;
  size_t stride = blockDim.x * gridDim.x;
  for (size_t i = index; i < n; i += stride) {
    Hmm[i] = Hmm[i] * 2 + 10;
  }
}


static size_t numElements{64 * 1024 * 1024};
static unsigned blocksPerCU{6};
static unsigned threadsPerBlock{256};

/*
   This testcase verifies the hipMallocManaged basic scenario - supported on all devices
 */
HIP_TEST_CASE(Unit_hipMallocManaged_Basic) {
  if (!HipTest::isManagedMemorySupportedOnDevice(0)) {
    WARN(
        "GPU doesn't support hipDeviceAttributeManagedMemory attribute so defaulting to system "
        "memory.");
  }

  float *A, *B, *C;

  HIP_CHECK(hipMallocManaged(&A, numElements * sizeof(float)));
  HIP_CHECK(hipMallocManaged(&B, numElements * sizeof(float)));
  HIP_CHECK(hipMallocManaged(&C, numElements * sizeof(float)));
  HIP_CHECK(hipFree(A));
  HIP_CHECK(hipFree(B));
  HIP_CHECK(hipFree(C));
}

/*
   This testcase verifies the hipMallocManaged advanced scenario - supported only on HMM enabled
   devices
 */
HIP_TEST_CASE(Unit_hipMallocManaged_Advanced) {
  CHECK_MANAGED_MEMORY_SUPPORT

  const size_t N = isQuickLevel() ? (512 * 1024 / sizeof(float)) : numElements;
  float *A, *B, *C;

  HIP_CHECK(hipMallocManaged(&A, N * sizeof(float)));
  HIP_CHECK(hipMallocManaged(&B, N * sizeof(float)));
  HIP_CHECK(hipMallocManaged(&C, N * sizeof(float)));
  HipTest::setDefaultData(N, A, B, C);

  hipDevice_t device = hipCpuDeviceId;

  HIP_CHECK(hipMemAdvise(A, N * sizeof(float), hipMemAdviseSetReadMostly, device));
  HIP_CHECK(hipMemPrefetchAsync(A, N * sizeof(float), 0));
  HIP_CHECK(hipMemPrefetchAsync(B, N * sizeof(float), 0));
  HIP_CHECK(hipDeviceSynchronize());
  HIP_CHECK(hipMemRangeGetAttribute(&device, sizeof(device),
                                    hipMemRangeAttributeLastPrefetchLocation, A,
                                    N * sizeof(float)));
  if (device != 0) {
    INFO("hipMemRangeGetAttribute error, device = " << device);
  }
  uint32_t read_only = 0xf;
  HIP_CHECK(hipMemRangeGetAttribute(&read_only, sizeof(read_only), hipMemRangeAttributeReadMostly,
                                    A, N * sizeof(float)));
  if (read_only != 1) {
    SUCCEED("hipMemRangeGetAttribute error, read_only = " << read_only);
  }

  unsigned blocks = HipTest::setNumBlocks(blocksPerCU, threadsPerBlock, N);
  hipEvent_t event0, event1;
  HIP_CHECK(hipEventCreate(&event0));
  HIP_CHECK(hipEventCreate(&event1));
  HIP_CHECK(hipEventRecord(event0, 0));
  hipLaunchKernelGGL(HipTest::vectorADD, dim3(blocks), dim3(threadsPerBlock), 0, 0,
                     static_cast<const float*>(A), static_cast<const float*>(B), C, N);
  HIP_CHECK(hipGetLastError());
  HIP_CHECK(hipEventRecord(event1, 0));
  HIP_CHECK(hipDeviceSynchronize());
  float time = 0.0f;
  HIP_CHECK(hipEventElapsedTime(&time, event0, event1));
  printf("Time %.3f ms\n", time);
  float maxError = 0.0f;
  HIP_CHECK(hipMemPrefetchAsync(B, N * sizeof(float), hipCpuDeviceId));
  HIP_CHECK(hipDeviceSynchronize());
  device = 0;
  HIP_CHECK(hipMemRangeGetAttribute(&device, sizeof(device),
                                    hipMemRangeAttributeLastPrefetchLocation, A,
                                    N * sizeof(float)));
  if (device != hipCpuDeviceId) {
    SUCCEED("hipMemRangeGetAttribute error device = " << device);
  }

  for (size_t i = 0; i < N; i++) {
    maxError = fmax(maxError, fabs(B[i] - 3.0f));
  }
  HIP_CHECK(hipFree(A));
  HIP_CHECK(hipFree(B));
  HIP_CHECK(hipFree(C));
  HIP_CHECK(hipEventDestroy(event0));
  HIP_CHECK(hipEventDestroy(event1));
  REQUIRE(maxError != 0.0f);
}

/*
   This testcase verifies that hipMallocManaged returns an OutOfMemory error for allocations much
   larger than the available memory - supported on all devices
 */
HIP_TEST_CASE(Unit_hipMallocManaged_Large) {
  if (!HipTest::isManagedMemorySupportedOnDevice(0)) {
    WARN(
        "GPU doesn't support hipDeviceAttributeManagedMemory attribute so defaulting to system "
        "memory.");
  }

  float* A;
  HIP_CHECK_ERROR(hipMallocManaged(&A, std::numeric_limits<size_t>::max()), hipErrorOutOfMemory);
}