// Copyright (c) 2021 CINN Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "paddle/cinn/backends/codegen_device_util.h"

#include "paddle/cinn/backends/cuda_util.h"
#include "paddle/cinn/ir/ir_mutator.h"
#include "paddle/cinn/optim/ir_simplify.h"
#include "paddle/common/enforce.h"

namespace cinn {
namespace backends {

std::tuple<ir::Module, ir::Module> SplitDeviceAndHostModule(ir::Module module) {
  detail::CollectBucketStrategyHostFunctionVisitor visitor(module->name);
  return visitor(module);
}

ir::Module CreateSwitchWithBroadcastConditionModule(
    const std::vector<ir::Expr> &broadcast_conditions,
    const std::vector<std::string> &case_func_names,
    const std::string &wrapper_func_name,
    const std::unordered_map<int, ir::Var> &symbolic_shape_var_index) {
  ir::Var kernel_args(KERNEL_ARGS, type_of<void *>());
  ir::Var kernel_args_num(KERNEL_ARGS_NUM, type_of<int>());
  ir::Var kernel_stream(KERNEL_STREAM, type_of<void *>());
  ir::Var tensor_shape_args(TENSOR_SHAPE_ARGS, type_of<int64_t **>());
  std::vector<ir::Argument> host_func_arguments = {
      ir::Argument(kernel_args, ir::Argument::IO::kOutput),
      ir::Argument(kernel_args_num, ir::Argument::IO::kInput),
      ir::Argument(kernel_stream, ir::Argument::IO::kOutput)};
  std::vector<ir::Argument> infer_shape_func_arguments = {
      ir::Argument(kernel_args, ir::Argument::IO::kOutput),
      ir::Argument(kernel_args_num, ir::Argument::IO::kInput),
      ir::Argument(tensor_shape_args, ir::Argument::IO::kOutput)};
  const auto &symbolic_arg_define = [&]() -> std::vector<ir::Expr> {
    std::vector<ir::Expr> arg_defs;
    for (const auto &item : symbolic_shape_var_index) {
      ir::Expr call_get_value_in_kernel_args =
          ir::Call::Make(Int(64),
                         runtime::intrinsic::get_value_in_kernel_args,
                         {kernel_args, ir::Expr(item.first)},
                         {},
                         ir::CallType::Extern,
                         ir::FunctionRef(),
                         0);
      ir::Expr let_symbol = ir::Expr(item.second);
      let_symbol->set_type(type_of<int64_t>());
      ir::Expr stmt = ir::Let::Make(let_symbol, call_get_value_in_kernel_args);
      arg_defs.push_back(stmt);
    }
    return arg_defs;
  }();

  const auto &CreateSwitchFunction =
      [&](std::vector<ir::Argument> func_arguments,
          const std::vector<ir::Expr> &read_args,
          std::string name_extend) -> ir::LoweredFunc {
    std::vector<ir::Expr> body_stmts(symbolic_arg_define);
    for (int i = 0; i < broadcast_conditions.size(); ++i) {
      ir::Expr callee = ir::Call::Make(Void(),
                                       case_func_names[i] + name_extend,
                                       read_args,
                                       {},
                                       ir::CallType::Extern,
                                       ir::FunctionRef(),
                                       0);
      if (i == 0) {
        body_stmts.emplace_back(
            ir::IfThenElse::Make(broadcast_conditions[i], callee));
      } else {
        auto false_expr = body_stmts.back();
        body_stmts.pop_back();
        body_stmts.emplace_back(
            ir::IfThenElse::Make(broadcast_conditions[i], callee, false_expr));
      }
    }
    ir::LoweredFunc caller =
        ir::_LoweredFunc_::Make(wrapper_func_name + name_extend,
                                func_arguments,
                                ir::Block::Make(body_stmts),
                                {});
    return caller;
  };

  ir::Module::Builder module_builder(wrapper_func_name + "_switch",
                                     cinn::common::DefaultHostTarget());
  ir::LoweredFunc host_func_caller = CreateSwitchFunction(
      host_func_arguments, {kernel_args, kernel_args_num, kernel_stream}, "");
  ir::LoweredFunc infer_shape_func_caller =
      CreateSwitchFunction(infer_shape_func_arguments,
                           {kernel_args, kernel_args_num, tensor_shape_args},
                           "_infer_shape");
  module_builder.AddFunctionWithoutOptim(host_func_caller);
  module_builder.AddFunctionWithoutOptim(infer_shape_func_caller);
  // no need cx86 func
  ir::LoweredFunc cx86_func_caller =
      ir::_LoweredFunc_::Make(wrapper_func_name + "_CX86",
                              host_func_arguments,
                              ir::Block::Make({}),
                              {});
  module_builder.AddFunctionWithoutOptim(cx86_func_caller);
  return module_builder.Build();
}

struct PredicatePrinter : public ir::IrPrinter {
  explicit PredicatePrinter(std::ostream &os) : ir::IrPrinter(os) {}

 private:
  void Visit(const ir::Add *x) { PrintBinaryOp("ADD", x); }
  void Visit(const ir::Sub *x) { PrintBinaryOp("SUB", x); }
  void Visit(const ir::Mul *x) { PrintBinaryOp("MUL", x); }
  void Visit(const ir::Div *x) { PrintBinaryOp("DIV", x); }
  void Visit(const ir::Mod *x) { PrintBinaryOp("MOD", x); }
  void Visit(const ir::EQ *x) { PrintBinaryOp("EQ", x); }
  void Visit(const ir::NE *x) { PrintBinaryOp("NE", x); }
  void Visit(const ir::LT *x) { PrintBinaryOp("LT", x); }
  void Visit(const ir::LE *x) { PrintBinaryOp("LE", x); }
  void Visit(const ir::GT *x) { PrintBinaryOp("GT", x); }
  void Visit(const ir::GE *x) { PrintBinaryOp("GE", x); }
  void Visit(const ir::And *x) { PrintBinaryOp("AND", x); }
  void Visit(const ir::Or *x) { PrintBinaryOp("OR", x); }
  void Visit(const ir::Max *x) { PrintBinaryOp("MAX", x); }
  void Visit(const ir::Min *x) { PrintBinaryOp("MIN", x); }
  void Visit(const ir::Call *x) { PrintCallOp(x); }

  template <typename IRN>
  void PrintBinaryOp(const std::string &op, const ir::BinaryOpNode<IRN> *x) {
    str_ += "_FPA_";
    ir::IrPrinter::Visit(x->a());
    str_ += op;
    ir::IrPrinter::Visit(x->b());
    str_ += "_BPA_";
  }

  void PrintCallOp(const ir::Call *x) {
    str_ += "_BCALL_";
    str_ += [&]() {
      std::string temp = x->name;
      std::transform(
          temp.begin(), temp.end(), temp.begin(), [](unsigned char c) {
            return std::toupper(c);
          });
      return temp;
    }();
    if (!x->read_args.empty()) {
      str_ += "_R_";
      for (const auto &v : x->read_args) ir::IrPrinter::Visit(v);
    }
    if (!x->write_args.empty()) {
      str_ += "_W_";
      for (const auto &v : x->write_args) ir::IrPrinter::Visit(v);
    }
    str_ += "_ECALL_";
  }
};

std::string Predicate2String(ir::Expr predicate) {
  std::stringstream ss;
  PredicatePrinter cond_printer(ss);
  cond_printer.Print(predicate);
  return ss.str();
}

static std::string CurTailFnName(const std::string &origin_fn_name) {
  const int MaxStrLength = 16383;
  if (origin_fn_name.length() <= MaxStrLength) {
    return origin_fn_name;
  }
  VLOG(6) << "Function name too long. Curtail and concat hash.";
  const std::string new_fn_name =
      origin_fn_name.substr(0, MaxStrLength) +
      std::to_string(std::hash<std::string>()(origin_fn_name));
  return new_fn_name;
}

bool RequiresCooperativeLaunch(const ir::LoweredFunc &func) {
  for (auto &space : func->temp_spaces) {
    if (space.size() != ir::Expr(0)) {
      return true;
    }
  }
  return false;
}

std::string
detail::CollectBucketStrategyHostFunctionVisitor::GenDeviceKernelName(
    const std::string &fn_name, ir::Expr predicate) {
  std::string cond_str = Predicate2String(predicate);
  // replace '-' with 'NEG'
  size_t pos = cond_str.find("-", 0);
  const std::string replacement_neg = "NEG";
  while (pos != std::string::npos) {
    cond_str.replace(pos, 1, replacement_neg);
    pos = cond_str.find("-", pos + replacement_neg.length());
  }

  // replace '!' with 'NOT'
  pos = cond_str.find("!", 0);
  const std::string replacement_not = "NOT";
  while (pos != std::string::npos) {
    cond_str.replace(pos, 1, replacement_not);
    pos = cond_str.find("!", pos + replacement_not.length());
  }
  VLOG(3) << "predicate string: " << cond_str;
  // NOTE(chenxi67): The kernel name is too long to be supported in cuda12.3 so
  // we need to curtail it.
  const std::string new_fn_name = CurTailFnName(fn_name);
  return new_fn_name + "_COND_" + cond_str + "__kernel";
}

void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
    ir::LoweredFunc func, ir::Expr predicate) {
  VLOG(4) << "Process Lowered Func" << func;
  ir::_LoweredFunc_ *func_node = func.As<ir::_LoweredFunc_>();
  PADDLE_ENFORCE_NOT_NULL(
      func_node,
      ::common::errors::InvalidArgument(
          "The provided function could not be cast to a lowered function. "
          "Please ensure the function is valid."));
  if (!func_node->cuda_axis_info.valid()) {
    func_node->cuda_axis_info.set_valid(true);
  }
  // process device func
  device_module_builder.AddFunctionWithoutOptim(
      CreateDeviceFunction(func, predicate));
  // process host func
  ir::Var kernel_ptr(GenDeviceKernelName(func_node->name, predicate),
                     type_of<std::string>());

  std::optional<Expr> shared_mem_bytes;
  cinn::common::DefaultDeviceTarget().arch.Match(
      [&](std::variant<common::UnknownArch, common::X86Arch, common::ARMArch>) {
        CINN_NOT_IMPLEMENTED;
      },
      [&](common::NVGPUArch) {
#ifdef CINN_WITH_CUDA
        shared_mem_bytes = CalculateSharedMemory(func);
#endif
      },
      [&](common::HygonDCUArchHIP) {
#ifdef CINN_WITH_HIP
        shared_mem_bytes = CalculateSharedMemory(func);
#endif
      },
      [&](common::HygonDCUArchSYCL) {
#ifdef CINN_WITH_SYCL
        shared_mem_bytes = Expr(0);
#endif
      },
      [&](common::CustomDeviceArch) {
#ifdef CINN_WITH_CUSTOM_DEVICE
        shared_mem_bytes = CalculateSharedMemory(func);
#endif
      });

  VLOG(6) << "Add a call node for func_node->name " << func_node->name << "\n"
          << "grid_dim: (" << func_node->cuda_axis_info.grid_dim(0) << ", "
          << func_node->cuda_axis_info.grid_dim(1) << ", "
          << func_node->cuda_axis_info.grid_dim(2) << "), "
          << "block_dim: (" << func_node->cuda_axis_info.block_dim(0) << ", "
          << func_node->cuda_axis_info.block_dim(1) << ", "
          << func_node->cuda_axis_info.block_dim(2) << "), "
          << "shared_mem: " << shared_mem_bytes.value();
  std::optional<const char *> call_kernel;
  cinn::common::DefaultDeviceTarget().arch.Match(
      [&](std::variant<common::UnknownArch, common::X86Arch, common::ARMArch>) {
        CINN_NOT_IMPLEMENTED;
      },
      [&](common::NVGPUArch) {
        call_kernel = RequiresCooperativeLaunch(func)
                          ? runtime::intrinsic::call_cuda_cooperative_kernel
                          : runtime::intrinsic::call_cuda_kernel;
      },
      [&](common::HygonDCUArchHIP) {
        call_kernel = runtime::intrinsic::call_hip_kernel;
      },
      [&](common::HygonDCUArchSYCL) {
        call_kernel = runtime::intrinsic::call_sycl_kernel;
      },
      [&](common::CustomDeviceArch) {
        call_kernel = runtime::intrinsic::call_custom_device_kernel;
      });
  // TODO(Dmovic): use new ir when backend update done.
  // Author(liujinnan): Copy args instead of use func args directly in host
  // func. because after longlong2int pass, some type of loweredfunc args may be
  // changed to int32, it cause compile error when lower to LLVM IR.
  std::vector<ir::Expr> kernel_args_int64 = {
      ir::ir_utils::IRCopy(func_node->cuda_axis_info.grid_dim(0)),
      ir::ir_utils::IRCopy(func_node->cuda_axis_info.grid_dim(1)),
      ir::ir_utils::IRCopy(func_node->cuda_axis_info.grid_dim(2)),
      ir::ir_utils::IRCopy(func_node->cuda_axis_info.block_dim(0)),
      ir::ir_utils::IRCopy(func_node->cuda_axis_info.block_dim(1)),
      ir::ir_utils::IRCopy(func_node->cuda_axis_info.block_dim(2)),
      ir::ir_utils::IRCopy(shared_mem_bytes.value()),
      cinn::common::make_const(Int(64), 0) /* enable TryElevateInt32ToInt64 */};
  ir::TryElevateInt32ToInt64(kernel_args_int64);

  ir::Expr call_extern_api =
      ir::Call::Make(Void(),
                     call_kernel.value(),
                     {kernel_ptr,
                      kernel_args_,
                      kernel_args_num_,
                      kernel_args_int64.at(0),  // grid_x
                      kernel_args_int64.at(1),  // grid_y
                      kernel_args_int64.at(2),  // grid_z
                      kernel_args_int64.at(3),  // block_x
                      kernel_args_int64.at(4),  // block_y
                      kernel_args_int64.at(5),  // block_z
                      kernel_args_int64.at(6),  // shared_mem
                      kernel_stream_},
                     {},
                     ir::CallType::Extern,
                     ir::FunctionRef(),
                     0);

  // create memset calls for temp_spaces if needed
  std::vector<ir::stmt::StmtRef> call_kernel_stmts;
  for (auto &temp_space : func_node->temp_spaces) {
    if (temp_space.need_zero_init()) {
      ir::Expr size = common::cast(temp_space.size(), common::UInt(64));
      ir::Expr call_get_arg =
          lang::CallExtern(runtime::intrinsic::get_item_in_cuda_kernel_args,
                           {kernel_args_, ir::Expr(temp_space.arg_idx())});
      ir::Expr call_memset = lang::CallExtern(
          runtime::intrinsic::call_cuda_memset,
          {call_get_arg, ir::Expr(1), ir::Expr(0), size, kernel_stream_});
      call_kernel_stmts.push_back(ir::stmt::Evaluate(call_memset));
    }
  }
  call_kernel_stmts.push_back(ir::stmt::Evaluate(call_extern_api));
  auto call_extern_api_block = ir::stmt::BlockRef(call_kernel_stmts);

  if (buckets_.empty()) {
    buckets_.emplace_back(
        ir::stmt::IfThenElse(predicate, call_extern_api_block));
  } else {
    auto false_expr = buckets_.back();
    buckets_.pop_back();
    buckets_.emplace_back(ir::stmt::IfThenElse(
        predicate,
        call_extern_api_block,
        ir::stmt::BlockRef(std::vector<ir::stmt::StmtRef>{false_expr})));
  }

  // create infer shape calls for temp_spaces
  std::vector<ir::stmt::StmtRef> temp_space_infer_shape_stmts;
  for (int i = 0; i < func_node->temp_spaces.size(); ++i) {
    ir::Var tensor_shape_args(TENSOR_SHAPE_ARGS, type_of<int64_t **>());
    ir::Expr size =
        common::cast(func_node->temp_spaces[i].size(), common::Int(64));
    ir::Expr call_set_value =
        lang::CallExtern(runtime::intrinsic::infer_shape_set_value,
                         {ir::Expr(func_node->num_output_tensors + i),
                          ir::Expr(0),
                          size,
                          tensor_shape_args});
    temp_space_infer_shape_stmts.push_back(ir::stmt::Evaluate(call_set_value));
  }
  if (!temp_space_infer_shape_stmts.empty()) {
    ir::stmt::BlockRef if_body =
        ir::stmt::BlockRef(temp_space_infer_shape_stmts);
    if (temp_space_infer_shape_body_.defined()) {
      temp_space_infer_shape_body_ = ir::stmt::IfThenElse(
          predicate,
          if_body,
          ir::stmt::BlockRef(
              std::vector<ir::stmt::StmtRef>{temp_space_infer_shape_body_}));
    } else {
      temp_space_infer_shape_body_ = ir::stmt::IfThenElse(predicate, if_body);
    }
  }
}

void detail::CollectBucketStrategyHostFunctionVisitor::ProcessArgs(
    ir::LoweredFunc func) {
  const std::vector<ir::Argument> &args = func->args;
  for (int i = 0; i < args.size(); ++i) {
    if (args[i].is_var()) {
      ir::Expr call_get_value_in_kernel_args =
          ir::Call::Make(Int(64),
                         runtime::intrinsic::get_value_in_kernel_args,
                         {kernel_args_, ir::Expr(i)},
                         {},
                         ir::CallType::Extern,
                         ir::FunctionRef(),
                         0);
      ir::Expr let_symbol = ir::ir_utils::IRCopy(args[i].var_arg());
      let_symbol->set_type(type_of<int64_t>());
      ir::stmt::StmtRef stmt =
          ir::stmt::Let(let_symbol, call_get_value_in_kernel_args);
      arg_defs_.push_back(stmt);
    }
  }
}

ir::LoweredFunc
detail::CollectBucketStrategyHostFunctionVisitor::CreateDeviceFunction(
    ir::LoweredFunc expr, ir::Expr predicate) {
  auto copied = ir::ir_utils::IRCopy(expr);
  copied->name = GenDeviceKernelName(copied->name, predicate);
  return copied;
}

}  // namespace backends
}  // namespace cinn