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

#include <hip_test_common.hh>

#include <hip/hiprtc.h>
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <cassert>
#include <cstddef>
#include <memory>
#include <iostream>
#include <iterator>
#include <vector>
#include "shfl.hh"

static constexpr auto shfl{
    R"(
template <typename T>
__global__ void shflUpSum(T* a, int size) {
  T val = a[threadIdx.x];
  auto all_threads = __activemask();
  for (unsigned int i = size / 2; i > 0; i /= 2) {
    val += __shfl_up_sync(all_threads, val, i, size);
  }
  a[threadIdx.x] = val;
}

template <typename T>
__global__ void shflDownSum(T* a, int size) {
  T val = a[threadIdx.x];
  auto all_threads = __activemask();
  for (int i = size / 2; i > 0; i /= 2) {
    val += __shfl_down_sync(all_threads, val, i, size);
  }
  a[threadIdx.x] = val;
}

template <typename T>
__global__ void shflXorSum(T* a, int size) {
  T val = a[threadIdx.x];
  auto all_threads = __activemask();
  for (int i = size/2; i > 0; i /= 2) {
    val += __shfl_xor_sync(all_threads, val, i, size);
  }
  a[threadIdx.x] = val;
}
)"};

template <typename T> void runTestShflSync(int option) {
  using namespace std;
  hiprtcProgram prog;
  hiprtcCreateProgram(&prog,           // prog
                      shfl,            // buffer
                      "shfl_sync.cu",  // name
                      0, nullptr, nullptr);

  string str;
  switch (option) {
    case 1:
      str = "shflUpSum<__half>";
      break;
    case 2:
      str = "shflDownSum<__half>";
      break;
    case 3:
      str = "shflXorSum<__half>";
      break;
    default:
      INFO("Options 1,2,3 are supported, but the passed option is: " << option);
      REQUIRE(false);
  }

  hiprtcAddNameExpression(prog, str.c_str());
  const char* options[] = {"-DHIP_ENABLE_WARP_SYNC_BUILTINS"};
  hiprtcResult compileResult{hiprtcCompileProgram(prog, 1, options)};
  size_t logSize;
  HIPRTC_CHECK(hiprtcGetProgramLogSize(prog, &logSize));
  if (logSize) {
    string log(logSize, '\0');
    HIPRTC_CHECK(hiprtcGetProgramLog(prog, &log[0]));
    std::cout << log << '\n';
  }
  REQUIRE(compileResult == HIPRTC_SUCCESS);
  size_t codeSize;
  HIPRTC_CHECK(hiprtcGetCodeSize(prog, &codeSize));

  vector<char> code(codeSize);
  HIPRTC_CHECK(hiprtcGetCode(prog, code.data()));

  // Do hip malloc first so that we donot need to do a cuInit manually before calling hipModule APIs
  size_t bufferSize = n * sizeof(T);

  T a[n];
  T cpuSum = sum(a);
  T* d_a;
  HIP_CHECK(hipMalloc(&d_a, bufferSize));

  hipModule_t module;
  hipFunction_t kernel;
  HIP_CHECK(hipModuleLoadData(&module, code.data()));
  const char* name;
  hiprtcGetLoweredName(prog, str.c_str(), &name);
  HIP_CHECK(hipModuleGetFunction(&kernel, module, name));

  HIP_CHECK(hipMemcpy(d_a, &a, bufferSize, hipMemcpyDefault));

  struct {
    T* a_;
    int b_;
  } args{d_a, n};

  auto size = sizeof(args);
  void* config[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, &args, HIP_LAUNCH_PARAM_BUFFER_SIZE, &size,
                    HIP_LAUNCH_PARAM_END};

  HIP_CHECK(hipModuleLaunchKernel(kernel, 1, 1, 1, n, 1, 1, 0, nullptr, nullptr, config));

  HIP_CHECK(hipMemcpy(&a, d_a, bufferSize, hipMemcpyDefault));
  bool result;
  switch (option) {
    case 1:  // shflUpSum
      result = compare(a[n - 1], cpuSum);
      break;
    case 2:  // shflDownSum
    case 3:  // shflXorSum
      result = compare(a[0], cpuSum);
      break;
  }

  if (result) {
    HIP_CHECK(hipFree(d_a));
    REQUIRE(false);
  }

  HIP_CHECK(hipFree(d_a));
  HIP_CHECK(hipModuleUnload(module));
  HIPRTC_CHECK(hiprtcDestroyProgram(&prog));
}

HIP_TEST_CASE(Unit_hiprtc_half_shuffle_sync) {
  runTestShflSync<__half>(1);
  runTestShflSync<__half>(2);
  runTestShflSync<__half>(3);
}