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

/**
 * @addtogroup hipMemcpy3D hipMemcpy3D
 * @{
 * @ingroup MemoryTest
 * `hipMemcpy3D(const hipMemcpy3DParms* p)` -
 * Copies data between 3D objects.
 */

/*
 * This testfile verifies the following scenarios of hipMemcpy3D API
 *
 * 1. Verifying hipMemcpy3D API for H2D,D2D and D2H scenarios for
      different datatypes and sizes.
 * 2. Verifying Negative Scenarios
 * 3. Verifying Extent validation scenarios by passing 0
 * 4. Verifying hipMemcpy3D API by allocating Memory in
 *    one GPU and trigger hipMemcpy3D from peer GPU
 *
 */

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

static constexpr auto width{10};
static constexpr auto height{10};
static constexpr auto depth{10};

template <typename T> class Memcpy3D {
  size_t width, height, depth;
  size_t size;
  hipArray_t arr, arr1;
  hipChannelFormatKind formatKind;
  hipMemcpy3DParms myparms;
  T* hData;

 public:
  Memcpy3D(size_t l_width, size_t l_height, size_t l_depth, hipChannelFormatKind l_format);
  void simple_Memcpy3D();
  void Extent_Validation();
  void NegativeTests();
  void AllocateMemory();
  void DeAllocateMemory();
  void SetDefaultData();
  void D2D_DeviceMem_OnDiffDevice();
  void D2H_H2D_DeviceMem_OnDiffDevice();
};

/*
 * This API sets the default values of hipMemcpy3DParms structure
 */
template <typename T> void Memcpy3D<T>::SetDefaultData() {
  myparms.srcPos = make_hipPos(0, 0, 0);
  myparms.dstPos = make_hipPos(0, 0, 0);
  myparms.extent = make_hipExtent(width, height, depth);
}

/*
 * Constructor initalized width,depth and height
 */
template <typename T> Memcpy3D<T>::Memcpy3D(size_t l_width, size_t l_height, size_t l_depth,
                                            hipChannelFormatKind l_format) {
  width = l_width;
  height = l_height;
  depth = l_depth;
  formatKind = l_format;
}

/*
 * Allocating Memory and initalizing data for both
 * device and host variables
 */
template <typename T> void Memcpy3D<T>::AllocateMemory() {
  size = width * height * depth * sizeof(T);
  hData = reinterpret_cast<T*>(malloc(size));
  memset(hData, 0, size);
  for (int i = 0; i < depth; i++) {
    for (int j = 0; j < height; j++) {
      for (int k = 0; k < width; k++) {
        hData[i * width * height + j * width + k] = i * width * height + j * width + k;
      }
    }
  }
  hipChannelFormatDesc channelDesc = hipCreateChannelDesc(sizeof(T) * 8, 0, 0, 0, formatKind);
  HIP_CHECK(
      hipMalloc3DArray(&arr, &channelDesc, make_hipExtent(width, height, depth), hipArrayDefault));
  HIP_CHECK(
      hipMalloc3DArray(&arr1, &channelDesc, make_hipExtent(width, height, depth), hipArrayDefault));
}

/*
 * DeAllocates the Memory of device and host variables
 */
template <typename T> void Memcpy3D<T>::DeAllocateMemory() {
  HIP_CHECK(hipFreeArray(arr));
  HIP_CHECK(hipFreeArray(arr1));
  free(hData);
}

/*
 * This API verifies both H2D & D2H functionalities of hipMemcpy3D API
 * by allocating memory in one GPU and calling the hipMemcpy3D API
 * from another GPU.
 * H2D case:
 * Input : "hData" is  initialized with the respective offset value
 * Output: Destination array "arr" variable.
 *
 * D2H case:
 * Input: "arr" array variable from the above output
 * Output: "hOutputData" variable data is copied from "arr" variable
 *
 * Validating the result by comparing "hData" and "hOutputData" variables
 */
template <typename T> void Memcpy3D<T>::D2H_H2D_DeviceMem_OnDiffDevice() {
  HIP_CHECK(hipSetDevice(0));
  int peerAccess = 0;
  HIP_CHECK(hipDeviceCanAccessPeer(&peerAccess, 1, 0));
  if (peerAccess) {
    AllocateMemory();
    // Memory is allocated on device 0 and Memcpy3DAsync triggered from device 1
    HIP_CHECK(hipSetDevice(1));

    // H2D Scenario
    memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
    SetDefaultData();
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.dstArray = arr;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyHostToDevice;
#else
    myparms.kind = hipMemcpyHostToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
    HIP_CHECK(hipDeviceSynchronize());

    // Device to host
    memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
    T* hOutputData = reinterpret_cast<T*>(malloc(size));
    memset(hOutputData, 0, size);
    SetDefaultData();
    myparms.dstPtr = make_hipPitchedPtr(hOutputData, width * sizeof(T), width, height);
    myparms.srcArray = arr;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToHost;
#else
    myparms.kind = hipMemcpyDeviceToHost;
#endif
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
    HIP_CHECK(hipDeviceSynchronize());

    // Validating the result
    HipTest::checkArray(hData, hOutputData, width, height, depth);
    free(hOutputData);
    DeAllocateMemory();
  } else {
    HIP_SKIP_TEST(HipTest::SkipReason::kPeerAccessUnavailable);
  }
}
/*
 * This API verifies both D2D functionalities of hipMemcpy3D API
 * by allocating memory in one GPU and calling the hipMemcpy3D API
 * from another GPU.
 *
 * D2D case:
 * Input : "arr" variable  is  initialized with the "hData" variable in GPU-0
 * Output: "arr2" variable in GPU-0
 *
 * hipMemcpy3D API is triggered from GPU-1
 * The "arr2" variable is then copied to "hOutputData" for validating
 * the result
 *
 * Validating the result by comparing "hData" and "hOutputData" variables
 */
template <typename T> void Memcpy3D<T>::D2D_DeviceMem_OnDiffDevice() {
  HIP_CHECK(hipSetDevice(0));
  int peerAccess = 0;
  HIP_CHECK(hipDeviceCanAccessPeer(&peerAccess, 0, 1));
  if (peerAccess) {
    AllocateMemory();
    memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
    SetDefaultData();

    // Host to device copy
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.dstArray = arr;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyHostToDevice;
#else
    myparms.kind = hipMemcpyHostToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
    hipArray_t arr2;
    hipChannelFormatDesc channelDesc1 = hipCreateChannelDesc(sizeof(T) * 8, 0, 0, 0, formatKind);
    HIP_CHECK(hipMalloc3DArray(&arr2, &channelDesc1, make_hipExtent(width, height, depth),
                               hipArrayDefault));

    // Allocating Mem on GPU device 0 and trigger hipMemcpy3D from GPU 1
    HIP_CHECK(hipSetDevice(1));

    // D2D Scenario
    memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
    SetDefaultData();
    myparms.srcArray = arr;
    myparms.dstArray = arr2;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToDevice;
#else
    myparms.kind = hipMemcpyDeviceToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
    HIP_CHECK(hipDeviceSynchronize());

    // For validating the D2D copy copying it again to hOutputData and
    // verifying it with iniital data hData
    memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
    T* hOutputData = reinterpret_cast<T*>(malloc(size));
    memset(hOutputData, 0, size);
    SetDefaultData();

    // Device to host
    myparms.dstPtr = make_hipPitchedPtr(hOutputData, width * sizeof(T), width, height);
    myparms.srcArray = arr2;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToHost;
#else
    myparms.kind = hipMemcpyDeviceToHost;
#endif
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
    HIP_CHECK(hipDeviceSynchronize());
    HipTest::checkArray(hData, hOutputData, width, height, depth);

    // DeAllocating Memory
    free(hOutputData);
    DeAllocateMemory();
  } else {
    HIP_SKIP_TEST(HipTest::SkipReason::kPeerAccessUnavailable);
  }
}
/*
 * This API verifies all the negative scenarios of hipMemcpy3D API
 */
template <typename T> void Memcpy3D<T>::NegativeTests() {
  HIP_CHECK(hipSetDevice(0));
  AllocateMemory();

  // Initialization of data
  memset(&myparms, 0, sizeof(myparms));
  myparms.srcPos = make_hipPos(0, 0, 0);
  myparms.dstPos = make_hipPos(0, 0, 0);
  myparms.extent = make_hipExtent(width, height, depth);
#ifdef __HIP_PLATFORM_NVIDIA__
  myparms.kind = cudaMemcpyHostToDevice;
#else
  myparms.kind = hipMemcpyHostToDevice;
#endif

  SECTION("Nullptr to destination array") {
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.dstArray = nullptr;
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Nullptr to source array") {
    myparms.srcArray = nullptr;
    myparms.dstPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing both Source ptr and array") {
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.srcArray = arr;
    myparms.dstArray = arr1;
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing both destination ptr and array") {
    myparms.dstPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.dstArray = arr;
    myparms.srcArray = arr1;
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing Max value to extent") {
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.dstArray = arr;
    myparms.extent =
        make_hipExtent(std::numeric_limits<int>::max(), std::numeric_limits<int>::max(),
                       std::numeric_limits<int>::max());
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing Source pitchedPtr as nullptr") {
    myparms.srcPtr = make_hipPitchedPtr(nullptr, width * sizeof(T), width, height);
    myparms.dstArray = arr;
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing Dst pitchedPtr as nullptr") {
    myparms.dstPtr = make_hipPitchedPtr(nullptr, width * sizeof(T), width, height);
    myparms.srcArray = arr;
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing width > max width size in extent") {
    myparms.extent = make_hipExtent(width + 1, height, depth);
    myparms.dstArray = arr;
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing hgt > max width size in extent") {
    myparms.extent = make_hipExtent(width, height + 1, depth);
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    myparms.dstArray = arr;
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing depth > max width size in extent") {
    myparms.extent = make_hipExtent(width, height, depth + 1);
    myparms.dstArray = arr;
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing dst width pos > max allocated width") {
    myparms.dstPos = make_hipPos(width + 1, 0, 0);
    myparms.dstArray = arr;
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing dst height pos > max allocated hgt") {
    myparms.dstPos = make_hipPos(0, height + 1, 0);
    myparms.dstArray = arr;
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing dst depth pos > max allocated depth") {
    myparms.dstPos = make_hipPos(0, 0, depth + 1);
    myparms.dstArray = arr;
    myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing src width pos > max allocated width") {
    myparms.srcPos = make_hipPos(width + 1, 0, 0);
    myparms.srcArray = arr;
    myparms.dstArray = arr1;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToDevice;
#else
    myparms.kind = hipMemcpyDeviceToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing src height pos > max allocated hgt") {
    myparms.srcPos = make_hipPos(0, height + 1, 0);
    myparms.srcArray = arr;
    myparms.dstArray = arr1;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToDevice;
#else
    myparms.kind = hipMemcpyDeviceToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing src height pos > max allocated hgt") {
    myparms.srcPos = make_hipPos(0, 0, depth + 1);
    myparms.srcArray = arr;
    myparms.dstArray = arr1;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToDevice;
#else
    myparms.kind = hipMemcpyDeviceToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  SECTION("Passing src array size  > dst array size") {
    // Passing src array size greater than destination array size
    hipArray_t arr2;
    hipChannelFormatDesc channelDesc1 = hipCreateChannelDesc(sizeof(T) * 8, 0, 0, 0, formatKind);
    HIP_CHECK(hipMalloc3DArray(&arr2, &channelDesc1, make_hipExtent(3, 3, 3), hipArrayDefault));
    myparms.srcArray = arr;
    myparms.dstArray = arr2;
#ifdef __HIP_PLATFORM_NVIDIA__
    myparms.kind = cudaMemcpyDeviceToDevice;
#else
    myparms.kind = hipMemcpyDeviceToDevice;
#endif
    REQUIRE(hipMemcpy3D(&myparms) != hipSuccess);
  }

  // DeAllocation of memory
  DeAllocateMemory();
}

/*
 * This API verifies the Extent validation Scenarios
 */
template <typename T> void Memcpy3D<T>::Extent_Validation() {
  HIP_CHECK(hipSetDevice(0));
  AllocateMemory();
  memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
  myparms.srcPos = make_hipPos(0, 0, 0);
  myparms.dstPos = make_hipPos(0, 0, 0);
  myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
  myparms.dstArray = arr;
#ifdef __HIP_PLATFORM_NVIDIA__
  myparms.kind = cudaMemcpyHostToDevice;
#else
  myparms.kind = hipMemcpyHostToDevice;
#endif
  SECTION("Passing Extent as 0") {
    myparms.extent = make_hipExtent(0, 0, 0);
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
  }
  SECTION("Passing Width 0 in Extent") {
    myparms.extent = make_hipExtent(0, height, depth);
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
  }
  SECTION("Passing Height 0 in Extent") {
    myparms.extent = make_hipExtent(width, 0, depth);
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
  }
  SECTION("Passing Depth 0 in Extent") {
    myparms.extent = make_hipExtent(width, height, 0);
    REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
  }
  SECTION("Passing Depth 0 in Extent") { REQUIRE(hipMemcpy3D(nullptr) != hipSuccess); }
  DeAllocateMemory();
}

/*
 * This API verifies H2H-D2D-D2H functionalities of hipMemcpy3D API
 *
 * Input : "arr" variable  is  initialized with the "hData" variable in GPU-0
 * Output: "arr1" variable in GPU-0
 *
 * The "arr1" variable is then copied to "hOutputData" for validating
 * the result
 *
 * Validating the result by comparing "hData" and "hOutputData" variables
 */

template <typename T> void Memcpy3D<T>::simple_Memcpy3D() {
  HIP_CHECK(hipSetDevice(0));
  AllocateMemory();

  // Host to Device
  memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
  SetDefaultData();
  myparms.srcPtr = make_hipPitchedPtr(hData, width * sizeof(T), width, height);
  myparms.dstArray = arr;
#ifdef __HIP_PLATFORM_NVIDIA__
  myparms.kind = cudaMemcpyHostToDevice;
#else
  myparms.kind = hipMemcpyHostToDevice;
#endif
  REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);

  // Array to Array
  memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
  SetDefaultData();
  myparms.srcArray = arr;
  myparms.dstArray = arr1;
#ifdef __HIP_PLATFORM_NVIDIA__
  myparms.kind = cudaMemcpyDeviceToDevice;
#else
  myparms.kind = hipMemcpyDeviceToDevice;
#endif
  REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);
  T* hOutputData = reinterpret_cast<T*>(malloc(size));
  memset(hOutputData, 0, size);

  // Device to host
  memset(&myparms, 0x0, sizeof(hipMemcpy3DParms));
  SetDefaultData();
  myparms.dstPtr = make_hipPitchedPtr(hOutputData, width * sizeof(T), width, height);
  myparms.srcArray = arr1;
#ifdef __HIP_PLATFORM_NVIDIA__
  myparms.kind = cudaMemcpyDeviceToHost;
#else
  myparms.kind = hipMemcpyDeviceToHost;
#endif
  REQUIRE(hipMemcpy3D(&myparms) == hipSuccess);

  // Validating the result
  HipTest::checkArray(hData, hOutputData, width, height, depth);

  // DeAllocating the Memory
  free(hOutputData);
  DeAllocateMemory();
}


/**
 * Test Description
 * ------------------------
 *  - This testcase performs the negative scenarios of hipMemcpy3D API
 * Test source
 * ------------------------
 *  - unit/memory/hipMemcpy3D_old.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */

HIP_TEST_CASE(Unit_hipMemcpy3D_multiDevice_Negative) {
  CHECK_IMAGE_SUPPORT
  int numDevices = 0;
  HIP_CHECK(hipGetDeviceCount(&numDevices));
  if (numDevices > 1) {
    Memcpy3D<int> memcpy3d(width, height, depth, hipChannelFormatKindSigned);
    memcpy3d.NegativeTests();
  } else {
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }
}

/**
 * End doxygen group MemoryTest.
 * @}
 */