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

/*
Added Negative and Functional tests for hipPointerGetAttribute API

Functional Scenarios:

1. Allocate memory using different Allocation APIs and check whether
   correct memory type and device oridinal are returned.

2. Allocate device variable and get the pointer info by calling hipPointerGetAttribute API
   with HIP_POINTER_ATTRIBUTE_DEVICE_POINTER/HIP_POINTER_ATTRIBUTE_START_ADDRESS
   and Launch kernel with device variable and verify whether the pointer variable of
   hipPointerGetAttribute is getting updated or not

3. Allocate device memory in GPU-0 and get the pointer info in peer GPU

4. Allocate device memory and get the buffer ID by calling
   hipPointerGetAttribute API with HIP_POINTER_ATTRIBUTE_BUFFER_ID,
   DeAllocate and Allocate the memory again and ensure that the buffer ID is unique

5. Allocate host memory and get the device ordinal by calling
   hipPointerGetAttribute API with HIP_POINTER_ATTRIBUTE_DEVICE_ORDINAL
   and ensure that it matches with CUDA result(which returns 100)

6. Allocate managed memory with different flags and trigger
   hipPointerGetAttribute with the following flags HIP_POINTER_ATTRIBUTE_MAPPED and verify the
behaviour
*/

#define checkVMMSupported(device) {                                                              \
  int value = 0;                                                                                 \
  hipDeviceAttribute_t attr = hipDeviceAttributeVirtualMemoryManagementSupported;                \
  HIP_CHECK(hipDeviceGetAttribute(&value, attr, device));                                        \
  if (value == 0) {                                                                              \
    HIP_SKIP_TEST(HipTest::SkipReason::kVmmUnsupported);                                         \
  }                                                                                              \
}

#include <hip_test_common.hh>
#include <string>
static constexpr auto NUM_W{16};
static constexpr auto NUM_H{16};
static constexpr size_t N{10};
#define INT_VAL 10
#define VAL_DATA 99
static __global__ void var_update(int* data) {
  for (unsigned int i = 0; i < N; i++) {
    data[i] = VAL_DATA;
  }
}

__managed__ int m_deferred_managed_var = 0;

/* Allocate memory using different Allocation APIs and check whether
   correct memory type and device oridinal are returned */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_MemoryTypes) {
  CHECK_IMAGE_SUPPORT

  HIP_CHECK(hipSetDevice(0));
  size_t pitch_A;
  size_t width{NUM_W * sizeof(char)};
  unsigned int datatype;
  SECTION("Malloc Pitch Allocation") {
    char* A_d;
    HIP_CHECK(hipMallocPitch(reinterpret_cast<void**>(&A_d), &pitch_A, width, NUM_H));
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                     reinterpret_cast<hipDeviceptr_t>(A_d)));

    REQUIRE(datatype == hipMemoryTypeDevice);
  }
#if HT_AMD
  SECTION("Malloc Array Allocation") {
    hipArray_t B_d;
    hipChannelFormatDesc desc = hipCreateChannelDesc<char>();
    HIP_CHECK(hipMallocArray(&B_d, &desc, NUM_W, NUM_H, hipArrayDefault));
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                     reinterpret_cast<hipDeviceptr_t>(B_d)));

    REQUIRE(datatype == hipMemoryTypeArray);
    HIP_CHECK(hipFreeArray(B_d));
  }

  SECTION("Malloc 3D Array Allocation") {
    int width = 10, height = 10, depth = 10;
    hipArray_t arr;

    hipChannelFormatDesc channelDesc =
        hipCreateChannelDesc(sizeof(float) * 8, 0, 0, 0, hipChannelFormatKindFloat);
    HIP_CHECK(hipMalloc3DArray(&arr, &channelDesc, make_hipExtent(width, height, depth),
                               hipArrayDefault));
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                     reinterpret_cast<hipDeviceptr_t>(arr)));

    REQUIRE(datatype == hipMemoryTypeArray);
    HIP_CHECK(hipFreeArray(arr));
  }
#endif
  SECTION("Unregistered Stack Pointer") {
    int stack_var = 0;
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                     reinterpret_cast<hipDeviceptr_t>(&stack_var)));
    REQUIRE(datatype == hipMemoryTypeUnregistered);
  }
  SECTION("Unregistered Heap Pointer") {
    int* heap_ptr = static_cast<int*>(malloc(sizeof(int)));
    REQUIRE(heap_ptr != nullptr);
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                     reinterpret_cast<hipDeviceptr_t>(heap_ptr)));
    REQUIRE(datatype == hipMemoryTypeUnregistered);
    free(heap_ptr);
  }
}

// A deferred managed variable must be classified as hipMemoryTypeManaged
// by hipPointerGetAttribute (HIP_POINTER_ATTRIBUTE_MEMORY_TYPE)
HIP_TEST_CASE(Unit_hipPointerGetAttribute_DeferredManagedVar) {
  CHECK_MANAGED_MEMORY_SUPPORT
  unsigned int datatype = 0;
  HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                   reinterpret_cast<hipDeviceptr_t>(&m_deferred_managed_var)));
  REQUIRE(datatype == hipMemoryTypeManaged);
}

/*
 * This testcase verifies the following scenario
 * Initializes A_d with A_h and get pointer info using hipPointerGetAttribute
 * The result of the API is passed to kernel for validation
 * and modifies it in kernel.
 * Validates the device variable to check whether the
 * data is updated or not.
 */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_KernelUpdation) {
  HIP_CHECK(hipSetDevice(0));
  size_t Nbytes = 0;
  Nbytes = N * sizeof(int);
  int *A_d, *A_h;
  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  hipDeviceptr_t data = 0;
  A_h = reinterpret_cast<int*>(malloc(Nbytes));
  for (unsigned int i = 0; i < N; i++) {
    A_h[i] = INT_VAL;
  }
  HIP_CHECK(hipMemcpy(A_d, A_h, Nbytes, hipMemcpyHostToDevice));
  HIP_CHECK(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_DEVICE_POINTER,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)));
  hipLaunchKernelGGL(var_update, dim3(1), dim3(1), 0, 0, reinterpret_cast<int*>(data));
  HIP_CHECK(hipGetLastError());
  HIP_CHECK(hipDeviceSynchronize());
  HIP_CHECK(hipMemcpy(A_h, A_d, Nbytes, hipMemcpyDeviceToHost));
  for (unsigned int i = 0; i < N; i++) {
    REQUIRE(A_h[i] == VAL_DATA);
  }
  HIP_CHECK(hipFree(A_d));
  free(A_h);
}
/*
 * This testcase verifies the pointer info of device variable
 * from peer GPU device.It validates the memory type and
 * device ordinal in peer GPU
 */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_PeerGPU) {
  HIP_CHECK(hipSetDevice(0));
  size_t Nbytes = 0;
  Nbytes = N * sizeof(int);
  int* A_d;
  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  unsigned int data = 0;
  int numDevices = 0;
  int canAccessPeer = 0;
  HIP_CHECK(hipGetDeviceCount(&numDevices));
  if (numDevices > 1) {
    HIP_CHECK(hipDeviceCanAccessPeer(&canAccessPeer, 0, 1));
    if (canAccessPeer) {
      HIP_CHECK(hipSetDevice(1));
      HIP_CHECK(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
                                       reinterpret_cast<hipDeviceptr_t>(A_d)));

      REQUIRE(data == hipMemoryTypeDevice);

      HIP_CHECK(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_DEVICE_ORDINAL,
                                       reinterpret_cast<hipDeviceptr_t>(A_d)));
      REQUIRE(data == 0);
    } else {
      HIP_SKIP_TEST(HipTest::SkipReason::kPeerAccessUnavailable);
    }
  } else {
    HIP_CHECK(hipFree(A_d));
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }
  HIP_CHECK(hipFree(A_d));
}

/* Allocate device memory and get the buffer ID by calling
   hipPointerGetAttribute API with HIP_POINTER_ATTRIBUTE_BUFFER_ID,
   DeAllocate and Allocate the memory again and
   ensure that the buffer ID is unique */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_BufferID) {
  HIP_CHECK(hipSetDevice(0));
  size_t Nbytes = 0;
  Nbytes = N * sizeof(int);
  int* A_d;
  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  unsigned int bufid1, bufid2;
  HIP_CHECK(hipPointerGetAttribute(&bufid1, HIP_POINTER_ATTRIBUTE_BUFFER_ID,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)));
  HIP_CHECK(hipFree(A_d));
  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  HIP_CHECK(hipPointerGetAttribute(&bufid2, HIP_POINTER_ATTRIBUTE_BUFFER_ID,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)));
  REQUIRE(bufid1 != bufid2);
  HIP_CHECK(hipFree(A_d));
}

/* Allocate host memory and get the device ordinal by calling
   hipPointerGetAttribute API with HIP_POINTER_ATTRIBUTE_DEVICE_ORDINAL
   and ensure that it matches with CUDA result
*/
#if HT_AMD
HIP_TEST_CASE(Unit_hipPointerGetAttribute_HostDeviceOrdinal) {
  size_t Nbytes = 0;
  Nbytes = N * sizeof(int);
  int* A_h;
  unsigned int data = 0, data1 = 0;
  A_h = reinterpret_cast<int*>(malloc(Nbytes));
  for (unsigned int i = 0; i < N; i++) {
    A_h[i] = INT_VAL;
  }
  REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_DEVICE_ORDINAL,
                                 reinterpret_cast<hipDeviceptr_t>(A_h)) == hipErrorInvalidValue);
  REQUIRE(hipPointerGetAttribute(&data1, HIP_POINTER_ATTRIBUTE_RANGE_SIZE,
                                 reinterpret_cast<hipDeviceptr_t>(A_h)) == hipErrorInvalidValue);
  free(A_h);
}
#endif

/* Allocate managed memory with different flags and trigger
   hipPointerGetAttribute with the following flags HIP_POINTER_ATTRIBUTE_MAPPED
   and verify the behaviour */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_MappedMem) {
  HIP_CHECK(hipSetDevice(0));
  size_t Nbytes = 0;
  Nbytes = N * sizeof(int);
  int *A_d, *A_h;
  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  A_h = reinterpret_cast<int*>(malloc(Nbytes));
  for (unsigned int i = 0; i < N; i++) {
    A_h[i] = INT_VAL;
  }
  HIP_CHECK(hipMemcpy(A_d, A_h, Nbytes, hipMemcpyHostToDevice));
  int *ptr1 = 0, *ptr2 = 0;
  unsigned int hostMalloc_mapped = 0;
  unsigned int mallocManaged = 0;
  HIP_CHECK(hipHostMalloc(&ptr1, Nbytes, hipHostMallocMapped));
  HIP_CHECK(hipMallocManaged(&ptr2, Nbytes, hipMemAttachGlobal));
  HIP_CHECK(hipPointerGetAttribute(&hostMalloc_mapped, HIP_POINTER_ATTRIBUTE_MAPPED,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)));
  HIP_CHECK(hipPointerGetAttribute(&mallocManaged, HIP_POINTER_ATTRIBUTE_MAPPED,
                                   reinterpret_cast<hipDeviceptr_t>(ptr2)));
  REQUIRE(hostMalloc_mapped == 1);
  REQUIRE(mallocManaged == 1);
  HIP_CHECK(hipFree(A_d));
  HIP_CHECK(hipHostFree(ptr1));
  HIP_CHECK(hipFree(ptr2));
  free(A_h);
}

/* This testcase verifies negative scenarios of hipPointerGetAttribute API */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_Negative) {
  HIP_CHECK(hipSetDevice(0));
  size_t Nbytes = 0;
  constexpr size_t N{100};
  Nbytes = N * sizeof(char);
  char* A_d;
  HIP_CHECK(hipMalloc(&A_d, Nbytes));
  hipDeviceptr_t data = 0;
  char* A_h;
  A_h = reinterpret_cast<char*>(malloc(Nbytes));
  SECTION("Pass nullptr to data") {
    REQUIRE(hipPointerGetAttribute(nullptr, HIP_POINTER_ATTRIBUTE_DEVICE_POINTER,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)) == hipErrorInvalidValue);
  }
  SECTION("Pass nullptr to device attribute") {
#if HT_AMD
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_DEVICE_POINTER, nullptr) ==
            hipErrorInvalidValue);
#else
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_DEVICE_POINTER,
                                   reinterpret_cast<hipDeviceptr_t>(nullptr)) ==
            hipErrorInvalidValue);
#endif
  }
#if !defined(ENABLE_ADDRESS_SANITIZER)
  SECTION("DeAllocateMem and get the pointer info") {
    char* B_d;
    HIP_CHECK(hipMalloc(&B_d, Nbytes));
    HIP_CHECK(hipFree(B_d));
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_DEVICE_POINTER,
                                   reinterpret_cast<hipDeviceptr_t>(B_d)) == hipErrorInvalidValue);
  }
#endif

  SECTION("Get Start address of host pointer") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR,
                                   reinterpret_cast<hipDeviceptr_t>(A_h)) == hipErrorInvalidValue);
  }
  SECTION("Pass HIP_POINTER_ATTRIBUTE_HOST_POINTER to device pointer") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_HOST_POINTER,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)) == hipErrorInvalidValue);
  }
  SECTION("Pass BUFFER_ID attribute to host pointer") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_BUFFER_ID,
                                   reinterpret_cast<hipDeviceptr_t>(A_h)) == hipErrorInvalidValue);
  }
  SECTION("Pass invalid attribute") {
    REQUIRE(hipPointerGetAttribute(&data, static_cast<hipPointer_attribute>(-1),
                                   reinterpret_cast<hipDeviceptr_t>(A_h)) == hipErrorInvalidValue);
  }
#if HT_AMD
  SECTION(
      "Pass HIP_POINTER_ATTRIBUTE_IS_GPU_DIRECT_RDMA_CAPABLE"
      "not supported by HIP") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_IS_GPU_DIRECT_RDMA_CAPABLE,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)) == hipErrorNotSupported);
  }
  SECTION("Pass HIP_POINTER_ATTRIBUTE_MEMPOOL_HANDLE not supported by HIP") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_MEMPOOL_HANDLE,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)) == hipErrorNotSupported);
  }
  SECTION("Pass HIP_POINTER_ATTRIBUTE_P2P_TOKENS  not supported by HIP") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_P2P_TOKENS,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)) == hipErrorNotSupported);
  }
  SECTION(
      "Pass HIP_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES"
      "not supported by HIP") {
    REQUIRE(hipPointerGetAttribute(&data, HIP_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES,
                                   reinterpret_cast<hipDeviceptr_t>(A_d)) == hipErrorNotSupported);
  }
#endif
  HIP_CHECK(hipFree(A_d));
  free(A_h);
}

/* Allocate memory using different Allocation APIs and check whether
   IPC CAPABLE attribute returns correctly */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_ipc_capable) {
  HIP_CHECK(hipSetDevice(0));
  size_t Nbytes = N * sizeof(int);
  unsigned int datatype;

  SECTION("Malloc Allocation") {
    int *A_d;
    HIP_CHECK(hipMalloc(&A_d, Nbytes));
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_IS_LEGACY_HIP_IPC_CAPABLE,
                                     reinterpret_cast<hipDeviceptr_t>(A_d)));
    REQUIRE(datatype == 1);
    HIP_CHECK(hipFree(A_d));
  }

  size_t pitch_A;
  size_t width{NUM_W * sizeof(char)};
  SECTION("Malloc Pitch Allocation") {
    char* A_d;
    HIP_CHECK(hipMallocPitch(reinterpret_cast<void**>(&A_d), &pitch_A, width, NUM_H));
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_IS_LEGACY_HIP_IPC_CAPABLE,
                                     reinterpret_cast<hipDeviceptr_t>(A_d)));

    REQUIRE(datatype == 1);
    HIP_CHECK(hipFree(A_d));
  }

  SECTION("VMM Memory Allocation") {
    size_t granularity = 0;
    int deviceId = 0;
    size_t buffer_size = N * sizeof(int);
    hipDevice_t device;
    HIP_CHECK(hipDeviceGet(&device, deviceId));
    checkVMMSupported(device);
    hipMemAllocationProp prop{};

    prop.type = hipMemAllocationTypePinned;
    prop.location.type = hipMemLocationTypeDevice;
    prop.location.id = device;  // Current Devices
    HIP_CHECK(
      hipMemGetAllocationGranularity(&granularity, &prop, hipMemAllocationGranularityMinimum));
    REQUIRE(granularity > 0);
    size_t size_mem = ((granularity + buffer_size - 1) / granularity) * granularity;
    hipMemGenericAllocationHandle_t handle;
    // Allocate physical memory
    HIP_CHECK(hipMemCreate(&handle, size_mem, &prop, 0));
    // Allocate virtual address range
    void* ptrA;
    HIP_CHECK(hipMemAddressReserve(&ptrA, size_mem, 0, 0, 0));
    HIP_CHECK(hipMemMap(ptrA, size_mem, 0, handle, 0));
    // Set access
    hipMemAccessDesc accessDesc = {};
    accessDesc.location.type = hipMemLocationTypeDevice;
    accessDesc.location.id = device;
    accessDesc.flags = hipMemAccessFlagsProtReadWrite;
    // Make the address accessible to GPU 0
    HIP_CHECK(hipMemSetAccess(ptrA, size_mem, &accessDesc, 1));
    HIP_CHECK(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_IS_LEGACY_HIP_IPC_CAPABLE,
                                     reinterpret_cast<hipDeviceptr_t>(ptrA)));

    REQUIRE(datatype == 0);
    HIP_CHECK(hipMemUnmap(ptrA, size_mem));
    HIP_CHECK(hipMemAddressFree(ptrA, size_mem));
    HIP_CHECK(hipMemRelease(handle));
 }

}

#if HT_AMD
/* HIP_POINTER_ATTRIBUTE_IS_LEGACY_HIP_IPC_CAPABLE on hipArray allocations (2D/3D). */
HIP_TEST_CASE(Unit_hipPointerGetAttribute_ipc_capable_Array) {
  CHECK_IMAGE_SUPPORT

  HIP_CHECK(hipSetDevice(0));
  unsigned int datatype;

  SECTION("Malloc Array Allocation") {
    hipArray_t B_d;
    hipChannelFormatDesc desc = hipCreateChannelDesc<char>();
    HIP_CHECK(hipMallocArray(&B_d, &desc, NUM_W, NUM_H, hipArrayDefault));
    HIP_CHECK_ERROR(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_IS_LEGACY_HIP_IPC_CAPABLE,
                                           reinterpret_cast<hipDeviceptr_t>(B_d)),
                    hipErrorInvalidValue);
    HIP_CHECK(hipFreeArray(B_d));
  }

  SECTION("Malloc 3D Array Allocation") {
    int width = 10, height = 10, depth = 10;
    hipArray_t arr;

    hipChannelFormatDesc channelDesc =
        hipCreateChannelDesc(sizeof(float) * 8, 0, 0, 0, hipChannelFormatKindFloat);
    HIP_CHECK(hipMalloc3DArray(&arr, &channelDesc, make_hipExtent(width, height, depth),
                               hipArrayDefault));
    HIP_CHECK_ERROR(hipPointerGetAttribute(&datatype, HIP_POINTER_ATTRIBUTE_IS_LEGACY_HIP_IPC_CAPABLE,
                                           reinterpret_cast<hipDeviceptr_t>(arr)),
                    hipErrorInvalidValue);
    HIP_CHECK(hipFreeArray(arr));
  }
}
#endif