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

#include <hip_test_common.hh>
#include <performance_common.hh>
#include <resource_guards.hh>

/**
 * @addtogroup stream stream
 * @{
 * @ingroup PerformanceTestStream
 * Contains performance tests for all hipStream related APIs
 */

class HipDeviceGetStreamPriorityRangeBenchmark
    : public Benchmark<HipDeviceGetStreamPriorityRangeBenchmark> {
 public:
  void operator()() {
    int priority_min, priority_max;
    TIMED_SECTION(kTimerTypeCpu) {
      HIP_CHECK(hipDeviceGetStreamPriorityRange(&priority_min, &priority_max));
    }
  }
};

class HipStreamQueryBenchmark : public Benchmark<HipStreamQueryBenchmark> {
 public:
  void operator()(bool perform_work) {
    hipError_t error;
    hipStream_t stream;
    HIP_CHECK(hipStreamCreate(&stream));
    void* dptr;

    if (perform_work) {
      HIP_CHECK(hipMallocAsync(&dptr, 2048 * 4, stream));
    }

    TIMED_SECTION(kTimerTypeCpu) { error = hipStreamQuery(stream); }

    if (perform_work) {
      HIP_CHECK(hipFreeAsync(dptr, stream));
      HIP_CHECK(hipStreamSynchronize(stream));
    }

    HIP_CHECK(hipStreamDestroy(stream));
  }
};

class HipStreamSynchronizeBenchmark : public Benchmark<HipStreamSynchronizeBenchmark> {
 public:
  void operator()() {
    hipError_t error;
    hipStream_t stream;
    HIP_CHECK(hipStreamCreate(&stream));

    TIMED_SECTION(kTimerTypeCpu) { error = hipStreamSynchronize(stream); }

    HIP_CHECK(hipStreamDestroy(stream));
  }
};

class HipStreamDestroyBenchmark : public Benchmark<HipStreamDestroyBenchmark> {
 public:
  void operator()() {
    hipStream_t stream;
    HIP_CHECK(hipStreamCreate(&stream));

    TIMED_SECTION(kTimerTypeCpu) { HIP_CHECK(hipStreamDestroy(stream)); }
  }
};

class HipStreamCreateBenchmark : public Benchmark<HipStreamCreateBenchmark> {
 public:
  void operator()() {
    hipStream_t stream;

    TIMED_SECTION(kTimerTypeCpu) { HIP_CHECK(hipStreamCreate(&stream)); }

    HIP_CHECK(hipStreamDestroy(stream));
  }
};

class HipStreamCreateWithPriorityBenchmark
    : public Benchmark<HipStreamCreateWithPriorityBenchmark> {
 public:
  void operator()(unsigned int flag) {
    hipStream_t stream;
    int priority_min, priority_max, priority_mid;

    HIP_CHECK(hipDeviceGetStreamPriorityRange(&priority_min, &priority_max));
    priority_mid = (priority_max + priority_min) / 2;

    TIMED_SECTION(kTimerTypeCpu) {
      HIP_CHECK(hipStreamCreateWithPriority(&stream, flag, priority_mid));
    }

    HIP_CHECK(hipStreamDestroy(stream));
  }
};


static std::string GetStreamCreateFlagName(unsigned flag) {
  switch (flag) {
    case hipStreamDefault:
      return "hipStreamDefault";
    case hipStreamNonBlocking:
      return "hipStreamNonBlocking";
    default:
      return "flag combination";
  }
}

class HipStreamCreateWithFlagsBenchmark : public Benchmark<HipStreamCreateWithFlagsBenchmark> {
 public:
  void operator()(unsigned int flag) {
    hipStream_t stream;

    TIMED_SECTION(kTimerTypeCpu) { HIP_CHECK(hipStreamCreateWithFlags(&stream, flag)); }

    HIP_CHECK(hipStreamDestroy(stream));
  }
};

/**
 * Test Description
 * ------------------------
 *  - Executes `hipStreamCreate`:
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipStreamCreate) {
  HipStreamCreateBenchmark benchmark;
  benchmark.Run();
}

static void RunBenchmark(unsigned flag) {
  HipStreamCreateWithFlagsBenchmark benchmark;
  benchmark.AddSectionName(GetStreamCreateFlagName(flag));
  benchmark.Run(flag);
}

static void RunBenchmarkWithPriority(unsigned flag) {
  HipStreamCreateWithPriorityBenchmark benchmark;
  benchmark.AddSectionName(GetStreamCreateFlagName(flag));
  benchmark.Run(flag);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipStreamCreateWithFlags` with all flags:
 *    -# Flags
 *      - hipStreamDefault
 *      - hipStreamNonBlocking
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipStreamCreateWithFlags) {
  const auto flag = GENERATE(hipStreamDefault, hipStreamNonBlocking);
  RunBenchmark(flag);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipStreamCreateWithPriority` with all flags:
 *    -# Flags
 *      - hipStreamDefault
 *      - hipStreamNonBlocking
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipStreamCreateWithPriority) {
  const auto flag = GENERATE(hipStreamDefault, hipStreamNonBlocking);
  RunBenchmarkWithPriority(flag);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipStreamDestroy`:
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipStreamDestroy) {
  HipStreamDestroyBenchmark benchmark;
  benchmark.Run();
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipDeviceGetStreamPriorityRange`:
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipDeviceGetStreamPriorityRange) {
  HipDeviceGetStreamPriorityRangeBenchmark benchmark;
  benchmark.Run();
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipStreamQuery`:
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipStreamQuery) {
  const auto perform_work = GENERATE(true, false);
  HipStreamQueryBenchmark benchmark;
  if (perform_work) {
    benchmark.AddSectionName("stream with work");
  } else {
    benchmark.AddSectionName("stream without work");
  }
  benchmark.Run(perform_work);
}

/**
 * Test Description
 * ------------------------
 *  - Executes `hipDeviceGetStreamPriorityRange`:
 * Test source
 * ------------------------
 * - performance/api/stream/hipStreamBasic.cc
 * Test requirements
 * ------------------------
 *  - HIP_VERSION >= 5.2
 */
HIP_TEST_CASE(Performance_hipStreamSynchronize) {
  HipStreamSynchronizeBenchmark benchmark;
  benchmark.Run();
}

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