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

#define HIP_ENABLE_WARP_SYNC_BUILTINS
#define HIP_ENABLE_EXTRA_WARP_SYNC_TYPES

#include "warp_common.hh"
#include <hip/hip_runtime.h>
#include <tuple>
#include <cmd_options.hh>
#include <functional>

#define NELEMS(array) (sizeof(array) / sizeof(array[0]))

// compiles the program, reusing the same compiling session for all the types
// (as opposed as calling the rtc compiler for each of the types)
template <template <typename> class Op, class T, typename... Types>
void compileProgram(hiprtcProgram& prog, const std::tuple<T, Types...>&) {
  std::string scalarName, intrinsicName, expression;
  std::tuple<Types...> remainingTypes;

  expression = std::string("reduceRtcKernel<") + typeToString<T>() + ", unsigned long long>";
  HIPRTC_CHECK(hiprtcAddNameExpression(prog, expression.c_str()));
  compileProgram<Op>(prog, remainingTypes);
}

template <class T, class MaskType, template <typename> class Op>
void runRtcReduceOp(hiprtcProgram& prog, T* output, const T* input, const MaskType* masks,
                    int numReduces, Op<T>) {
  unsigned int wavefrontSize = getWarpSize();
  const char* loweredName;
  hipFunction_t kernel;
  hipModule_t module;
  LinearAllocGuard<int> d_numReduces(LinearAllocs::hipMalloc, sizeof(int));

  HIP_CHECK(hipMemcpy(d_numReduces.ptr(), &numReduces, sizeof(int), hipMemcpyHostToDevice));
  std::vector<const void*> args = {output, input, masks, d_numReduces.ptr()};
  std::size_t sizeBytes = args.size() * sizeof(void*);
  void* config[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, args.data(), HIP_LAUNCH_PARAM_BUFFER_SIZE,
                    &sizeBytes, HIP_LAUNCH_PARAM_END};
  std::vector<char> code;
  size_t codeSize;
  std::string expression =
      std::string("reduceRtcKernel<") + typeToString<T>() + ", unsigned long long>";
  dim3 grdDim{1u};
  dim3 blkDim{wavefrontSize};

  HIPRTC_CHECK(hiprtcGetCodeSize(prog, &codeSize));
  code.resize(codeSize);
  HIPRTC_CHECK(hiprtcGetCode(prog, code.data()));
  HIP_CHECK(hipModuleLoadData(&module, code.data()));
  HIPRTC_CHECK(hiprtcGetLoweredName(prog, expression.c_str(), &loweredName));
  HIP_CHECK(hipModuleGetFunction(&kernel, module, loweredName));
  HIP_CHECK(hipModuleLaunchKernel(kernel, grdDim.x, grdDim.y, grdDim.z, blkDim.x, blkDim.y,
                                  blkDim.z, 0, 0, nullptr, config));
  HIP_CHECK(hipModuleUnload(module));
}

template <template <typename> class Op, class Type = void>
void runTestReduceForTypes(hiprtcProgram&, const std::tuple<>) {}

template <template <typename> class Op, class T, typename... Types>
void runTestReduceForTypes(hiprtcProgram& prog, const std::tuple<T, Types...>) {
  std::tuple<Types...> remainingTypes;
  int iteration = 0;

  auto reduceFunc = [&prog](T* d_output, const T* d_input, const unsigned long long* d_masks,
                            int numReduces, Op<T> op) {
    runRtcReduceOp(prog, d_output, d_input, d_masks, numReduces, op);
  };

  while (iteration < cmd_options.reduce_iterations) {
    runTestReduce<T, decltype(reduceFunc), Op>(iteration, reduceFunc);
    iteration++;

    if (cmd_options.reduce_iterations != 1) {
      std::cout << "\rIteration: " << iteration;
      std::flush(std::cout);
    }
  }

  runTestReduceForTypes<Op>(prog, remainingTypes);
}

template <class T, template <typename> class Op>
void reduceOpToString(std::string& scalarName, std::string& intrinsicName) {
  if constexpr (std::is_same<Op<T>, std::plus<T>>::value) {
    scalarName = "std::plus";
    intrinsicName = "__reduce_add_sync";
  } else if constexpr (std::is_same<Op<T>, MinOp<T>>::value) {
    scalarName = "MinOp";
    intrinsicName = "__reduce_min_sync";
  } else if constexpr (std::is_same<Op<T>, MaxOp<T>>::value) {
    scalarName = "MaxOp";
    intrinsicName = "__reduce_max_sync";
  } else if constexpr (std::is_same<Op<T>, AndOp<T>>::value) {
    scalarName = "std::bit_and";
    intrinsicName = "__reduce_and_sync";
  } else if constexpr (std::is_same<Op<T>, OrOp<T>>::value) {
    scalarName = "std::bit_or";
    intrinsicName = "__reduce_or_sync";
  } else if constexpr (std::is_same<Op<T>, XorOp<T>>::value) {
    scalarName = "std::bit_xor";
    intrinsicName = "__reduce_xor_sync";
  } else
    static_assert(std::is_void<T>::value, "Unexpected operator");
}

template <template <typename> class Op, class T = void>
void compileProgram(hiprtcProgram& prog, const std::tuple<>&) {
  size_t logSize;
  std::string scalarName, intrinsicName;
  hiprtcResult compileResult;
  const char* options[] = {"-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-DHIP_ENABLE_EXTRA_WARP_SYNC_TYPES"};

  reduceOpToString<int, Op>(scalarName, intrinsicName);
  compileResult = hiprtcResult{hiprtcCompileProgram(prog, NELEMS(options), options)};
  HIPRTC_CHECK(hiprtcGetProgramLogSize(prog, &logSize));

  if (compileResult != HIPRTC_SUCCESS || logSize > 0) {
    std::string log(logSize, '\0');

    HIPRTC_CHECK(hiprtcGetProgramLog(prog, &log[0]));
    std::cerr << "Runtime compilation failed or contained warnings for operator: " << scalarName
              << " associated reduce function: " << intrinsicName << "\n";
    std::cerr << log << '\n';
    REQUIRE(false);
  }
}

template <template <typename> class Op, typename... Types>
void runAndCompileTest(const std::tuple<Types...> types) {
  std::string scalarName, intrinsicName, kernelStr;
  hiprtcProgram prog;

  reduceOpToString<int, Op>(scalarName, intrinsicName);
  kernelStr = R"(
    template <class T, class MaskType>
    __global__ void reduceRtcKernel(T* output, const T* input, const MaskType* masks, int* numReduces)
    {
      int tid = threadIdx.x;
      int laneId = tid % warpSize;

      for (int i = 0; i < *numReduces; i++) {
        int idx = warpSize * i + laneId;
        if (masks[i] & (1ul << tid)) {
          // call the operator only if the lane is mentioned in the mask
          T& result = output[idx];
          result = )" +
              intrinsicName + R"((masks[i], input[idx]);
        }
      }
   })";

  HIPRTC_CHECK(
      hiprtcCreateProgram(&prog, kernelStr.c_str(), "warp_reduce.hip", 0, nullptr, nullptr));
  compileProgram<Op>(prog, types);
  runTestReduceForTypes<Op>(prog, types);
  HIPRTC_CHECK(hiprtcDestroyProgram(&prog));
}

HIP_TEST_CASE(Unit_Rtc_ReduceRandom) {
  const std::tuple<int, unsigned int, long long, unsigned long long, float, half, double> allTypes;
  const std::tuple<int, unsigned int, long long, unsigned long long> integralTypes;

  SECTION("add") { runAndCompileTest<std::plus>(allTypes); }

  SECTION("min") { runAndCompileTest<MinOp>(allTypes); }

  SECTION("max") { runAndCompileTest<MaxOp>(allTypes); }

  SECTION("and") { runAndCompileTest<AndOp>(integralTypes); }

  SECTION("or") { runAndCompileTest<OrOp>(integralTypes); }

  SECTION("xor") { runAndCompileTest<XorOp>(integralTypes); }
}