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

#include <hip_test_common.hh>

#include <algorithm>
#include <atomic>
#include <map>
#include <mutex>
#include <random>
#include <thread>

__global__ void addVal(unsigned long long* ptr, size_t index, unsigned long long val) {
  atomicAdd(ptr + index, val);
}

// Create a copy constructible AtomicWrap around std::atomic so that
// we can put it in a vector
template <typename T> struct AtomicWrap {
  std::atomic<T> data;

  AtomicWrap() : data() {}

  AtomicWrap(T i) : data(i) {}

  AtomicWrap(const std::atomic<T>& a) : data(a.load()) {}

  AtomicWrap(const AtomicWrap& other) : data(other.data.load()) {}

  AtomicWrap& operator=(const AtomicWrap& other) {
    data.store(other.data.load());
    return *this;
  }
};

// Have multiple threads and enqueue commands from them on a single stream
// Validate at the end that all commands have completed successfully
HIP_TEST_CASE(Stress_StreamEnqueue_DifferentThreads) {
  auto hwThreads = std::thread::hardware_concurrency();
  hwThreads = (hwThreads >= 2) ? hwThreads : 2;  // Run atleast 2 threads

  std::vector<AtomicWrap<unsigned long long>> hostData(hwThreads, 0);

  unsigned long long* dPtr{nullptr};
  HIP_CHECK(hipMalloc(&dPtr, sizeof(unsigned long long) * hwThreads));
  REQUIRE(dPtr != nullptr);

  HIP_CHECK(hipMemset(dPtr, 0, sizeof(unsigned long long) * hwThreads));

  std::random_device device;
  std::mt19937 engine(device());

  constexpr size_t maxWork = 10000;
  constexpr size_t maxVal = 10;

  std::uniform_int_distribution<std::mt19937::result_type> genIndex(0, hwThreads - 1);
  std::uniform_int_distribution<std::mt19937::result_type> genWork(0, maxWork);
  std::uniform_int_distribution<std::mt19937::result_type> genVal(0, maxVal);

  auto enqueueKernelThread = [&](hipStream_t stream) {
    auto iter = genWork(engine);  // Generate work to be done via thread
    for (unsigned long i = 0; i < iter; i++) {
      auto index = genIndex(engine);  // Generate Index to add to
      auto val = genVal(engine);      // Generate value to add to the destination
      hostData[index].data += val;    // Replicate it on host
      addVal<<<1, 1, 0, stream>>>(dPtr, static_cast<size_t>(index),
                                  static_cast<unsigned long long>(val));  // And on device
    }
  };

  hipStream_t stream{};
  HIP_CHECK(hipStreamCreate(&stream));
  std::vector<std::thread> threadPool{};
  threadPool.reserve(hwThreads);

  // Launch work
  for (size_t i = 0; i < hwThreads; i++) {
    threadPool.emplace_back(std::thread(enqueueKernelThread, stream));
  }

  // Wait for work to finish
  for (auto& i : threadPool) {
    i.join();
  }

  HIP_CHECK(hipStreamDestroy(stream));

  auto hPtr = std::make_unique<unsigned long long[]>(hwThreads);
  HIP_CHECK(
      hipMemcpy(hPtr.get(), dPtr, sizeof(unsigned long long) * hwThreads, hipMemcpyDeviceToHost));

  HIP_CHECK(hipFree(dPtr));

  // Validate that CPU and GPU has the same results
  for (size_t i = 0; i < hwThreads; i++) {
    INFO("Check for Index " << i);
    REQUIRE(hostData[i].data.load() == hPtr[i]);
  }
}

__global__ void doOperation(int* dPtr, int val) {
  auto i = threadIdx.x;
  atomicAdd(dPtr + i, val);
}

// Allocate mulitple stream for same device.
// Same device stream operate on same memory
HIP_TEST_CASE(Stress_StreamEnqueue_DifferentThreads_MultiGPU) {
  int deviceCount{0};
  HIP_CHECK(hipGetDeviceCount(&deviceCount));
  REQUIRE(deviceCount > 0);

  // Skip the test if devices less than 2
  if (deviceCount <= 1) {
    HIP_SKIP_TEST(HipTest::SkipReason::kFewerThanTwoGpus);
  }

  constexpr size_t streamPerGPU{3};  // Stream per gpu

  std::vector<hipStream_t> streamPool{};
  streamPool.reserve(deviceCount * streamPerGPU);
  // Map of stream and device memory
  std::map<hipStream_t, int*> streamToDeviceMemory;
  // Map of stream and host result
  std::map<hipStream_t, AtomicWrap<int>> streamToHostMemory;
  // Map of stream and device it was created on
  std::map<hipStream_t, size_t> streamToDeviceIndex;
  constexpr size_t size = 1024;

  for (int i = 0; i < deviceCount; i++) {
    HIP_CHECK(hipSetDevice(i));

    for (size_t j = 0; j < streamPerGPU; j++) {
      hipStream_t stream{nullptr};
      HIP_CHECK(hipStreamCreate(&stream));
      REQUIRE(stream != nullptr);
      streamPool.push_back(stream);

      int* dPtr{nullptr};
      HIP_CHECK(hipMalloc(&dPtr, sizeof(int) * size));
      REQUIRE(dPtr != nullptr);
      HIP_CHECK(hipMemset(dPtr, 0, sizeof(int) * size));
      // All streams work on exclusive memory
      streamToDeviceMemory[stream] = dPtr;

      streamToHostMemory[stream] = AtomicWrap<int>(0);  // CPU result

      streamToDeviceIndex[stream] = i;  // Capture device id for stream
    }
  }

  constexpr size_t maxVal = 5;
  constexpr size_t maxWorkPerThread = 10000;

  // Boiler plate code to generate a random number
  std::random_device device;
  std::mt19937 engine(device());

  std::uniform_int_distribution<std::mt19937::result_type> genVal(-maxVal, maxVal);
  std::uniform_int_distribution<std::mt19937::result_type> genStream(0, streamPool.size() - 1);

#if HT_NVIDIA
  std::mutex ness;  // On nvidia, current device needs to match stream's device
#endif

  auto enqueueKernelThread = [&]() {
    for (size_t i = 0; i < maxWorkPerThread; i++) {
#if HT_NVIDIA
      std::unique_lock<std::mutex> lock(ness);  // Lock on creation
#endif
      // Get a random stream
      hipStream_t stream = streamPool[genStream(engine)];

      // TODO use HIP_CHECK_THREAD when PR#2664 is merged
      if (hipSuccess != hipSetDevice(streamToDeviceIndex[stream])) {
        return;
      }

      int val = genVal(engine);  // Generate Value to add/sub to
      // Replicate result on CPU
      streamToHostMemory[stream].data.fetch_add(val);
      auto dPtr = streamToDeviceMemory[stream];
      doOperation<<<1, 1024, 0, stream>>>(dPtr, val);  // On GPU
    }
  };

  auto maxThreads = std::thread::hardware_concurrency();
  maxThreads = (maxThreads >= 2) ? maxThreads : 2;  // Run atleast 2 threads

  std::vector<std::thread> threadPool{};
  threadPool.reserve(maxThreads);

  // Launch Threads
  for (size_t i = 0; i < maxThreads; i++) {
    threadPool.emplace_back(std::thread(enqueueKernelThread));
  }

  // Wait for them to stop
  for (auto& i : threadPool) {
    i.join();
  }

  // Sync and check results
  for (auto& i : streamPool) {
    HIP_CHECK(hipStreamSynchronize(i));
    auto dResult = std::make_unique<int[]>(size);
    HIP_CHECK(hipMemcpy(dResult.get(), streamToDeviceMemory[i], sizeof(int) * size,
                        hipMemcpyDeviceToHost));
    HIP_CHECK(hipFree(streamToDeviceMemory[i]));
    HIP_CHECK(hipStreamDestroy(i));
    auto res = streamToHostMemory[i].data.load();
    INFO("Matching CPU: " << res << " GPU: " << dResult[0] << " Dev Ptr: "
                          << streamToDeviceMemory[i] << " on Device: " << streamToDeviceIndex[i]);
    REQUIRE(std::all_of(dResult.get(), dResult.get() + size, [=](int r) { return r == res; }));
  }
}