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

#include "memcpy_performance_common.hh"

/**
 * @addtogroup memcpy memcpy
 * @{
 * @ingroup PerformanceTestMemory
 */

class Memcpy2DBenchmark : public Benchmark<Memcpy2DBenchmark> {
 public:
  void operator()(void* dst, size_t dst_pitch, const void* src, size_t src_pitch, size_t width,
                  size_t height, hipMemcpyKind kind) {
    TIMED_SECTION(kTimerTypeCpu) {
      HIP_CHECK(hipMemcpy2D(dst, dst_pitch, src, src_pitch, width, height, kind));
    }
  }
};

static void RunBenchmark(size_t width, size_t height, hipMemcpyKind kind,
                         bool enable_peer_access = false) {
  Memcpy2DBenchmark benchmark;
  benchmark.AddSectionName("(" + std::to_string(width) + ", " + std::to_string(height) + ")");

  if (kind == hipMemcpyDeviceToHost) {
    LinearAllocGuard2D<int> device_allocation(width, height);
    LinearAllocGuard<int> host_allocation(LinearAllocs::hipHostMalloc,
                                          device_allocation.width() * height);
    benchmark.Run(host_allocation.ptr(), device_allocation.width(), device_allocation.ptr(),
                  device_allocation.pitch(), device_allocation.width(), device_allocation.height(),
                  hipMemcpyDeviceToHost);
  } else if (kind == hipMemcpyHostToDevice) {
    LinearAllocGuard2D<int> device_allocation(width, height);
    LinearAllocGuard<int> host_allocation(LinearAllocs::hipHostMalloc,
                                          device_allocation.width() * height);
    benchmark.Run(device_allocation.ptr(), device_allocation.pitch(), host_allocation.ptr(),
                  device_allocation.width(), device_allocation.width(), device_allocation.height(),
                  hipMemcpyHostToDevice);
  } else if (kind == hipMemcpyHostToHost) {
    LinearAllocGuard<int> src_allocation(LinearAllocs::hipHostMalloc, width * sizeof(int) * height);
    LinearAllocGuard<int> dst_allocation(LinearAllocs::hipHostMalloc, width * sizeof(int) * height);
    benchmark.Run(dst_allocation.ptr(), width * sizeof(int), src_allocation.ptr(),
                  width * sizeof(int), width * sizeof(int), height, hipMemcpyHostToHost);
  } else {
    // hipMemcpyDeviceToDevice
    int src_device = std::get<0>(GetDeviceIds(enable_peer_access));
    int dst_device = std::get<1>(GetDeviceIds(enable_peer_access));

    LinearAllocGuard2D<int> src_allocation(width, height);
    HIP_CHECK(hipSetDevice(dst_device));
    LinearAllocGuard2D<int> dst_allocation(width, height);
    HIP_CHECK(hipSetDevice(src_device));
    benchmark.Run(dst_allocation.ptr(), dst_allocation.pitch(), src_allocation.ptr(),
                  src_allocation.pitch(), dst_allocation.width(), dst_allocation.height(),
                  hipMemcpyDeviceToDevice);
  }
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipMemcpy2D` from Device to Host:
 *    -# Allocation size
 *      - Small: 4 KB x 32 B
 *      - Medium: 4 MB x 32 B
 *      - Large: 16 MB x 32 B
 * Test source
 * ------------------------
 * - performance/api/memcpy/hipMemcpy2D.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipMemcpy2D_DeviceToHost) {
  const auto width = GENERATE(4_KB, 4_MB, 16_MB);
  RunBenchmark(width, 32, hipMemcpyDeviceToHost);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipMemcpy2D` from Host to Device:
 *    -# Allocation size
 *      - Small: 4 KB x 32 B
 *      - Medium: 4 MB x 32 B
 *      - Large: 16 MB x 32 B
 * Test source
 * ------------------------
 * - performance/api/memcpy/hipMemcpy2D.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipMemcpy2D_HostToDevice) {
  const auto width = GENERATE(4_KB, 4_MB, 16_MB);
  RunBenchmark(width, 32, hipMemcpyHostToDevice);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipMemcpy2D` from Host to Host:
 *    -# Allocation size
 *      - Small: 4 KB x 32 B
 *      - Medium: 4 MB x 32 B
 *      - Large: 16 MB x 32 B
 * Test source
 * ------------------------
 * - performance/api/memcpy/hipMemcpy2D.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipMemcpy2D_HostToHost) {
  const auto width = GENERATE(4_KB, 4_MB, 16_MB);
  RunBenchmark(width, 32, hipMemcpyHostToHost);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipMemcpy2D` from Device to Device with peer access disabled:
 *    -# Allocation size
 *      - Small: 4 KB x 32 B
 *      - Medium: 4 MB x 32 B
 *      - Large: 16 MB x 32 B
 * Test source
 * ------------------------
 * - performance/api/memcpy/hipMemcpy2D.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipMemcpy2D_DeviceToDevice_DisablePeerAccess) {
  const auto width = GENERATE(4_KB, 4_MB, 16_MB);
  RunBenchmark(width, 32, hipMemcpyDeviceToDevice);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipMemcpy2D` from Device to Device with peer access enabled:
 *    -# Allocation size
 *      - Small: 4 KB x 32 B
 *      - Medium: 4 MB x 32 B
 *      - Large: 16 MB x 32 B
 * Test source
 * ------------------------
 * - performance/api/memcpy/hipMemcpy2D.cc
 * Test requirements
 * ------------------------
 *  - Multi-device
 *  - Device supports Peer-to-Peer access
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipMemcpy2D_DeviceToDevice_EnablePeerAccess) {
  if (HipTest::getDeviceCount() < 2) {
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }
  const auto width = GENERATE(4_KB, 4_MB, 16_MB);
  RunBenchmark(width, 32, hipMemcpyDeviceToDevice, true);
}

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