/*========================== begin_copyright_notice ============================

Copyright (C) 2024 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#include "Compiler/Optimizer/OpenCLPasses/PoisonFP64KernelsPass/PoisonFP64KernelsPass.h"
#include "Compiler/Optimizer/OpenCLPasses/KernelArgs/KernelArgs.hpp"
#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include <llvm/Pass.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/Analysis/CallGraph.h>
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;
using namespace IGC;

char PoisonFP64KernelsLPM::ID = 0;
const char *PoisonFP64Kernels::attributeName = "invalid_kernel(\"uses-fp64-math\")";
// Register pass to igc-opt
#define PASS_FLAG "igc-poison-fp64-kernels"
#define PASS_DESCRIPTION "Poison kernels using FP64 on unsupported platforms."
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(PoisonFP64KernelsLPM, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_END(PoisonFP64KernelsLPM, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

PoisonFP64KernelsLPM::PoisonFP64KernelsLPM() : CallGraphSCCPass(ID) {
  initializePoisonFP64KernelsLPMPass(*PassRegistry::getPassRegistry());
}

void PoisonFP64Kernels::markForRemoval(Function *F) {
  if (fp64Functions.count(F) > 0)
    return;

  fp64FunctionsOrder.push_back(F);
  fp64Functions.insert(F);
}

void PoisonFP64Kernels::reset() {
  fp64Functions.clear();
  fp64FunctionsOrder.clear();
}

bool PoisonFP64Kernels::runOnSCCNodes(const std::vector<CallGraphNode *> &SCC, CodeGenContext *ctx) {
  // If DP emu is enabled we don't need poison fp64 pass
  if (ctx->m_hasDPEmu) {
    return false;
  }

  bool modified = false;
  for (CallGraphNode *Node : SCC) {
    Function *F = Node->getFunction();
    if (F == nullptr || F->isDeclaration())
      continue;

    bool markedCallee = false;
    for (BasicBlock &BB : *F) {
      for (Instruction &I : BB) {
        auto *CB = dyn_cast<CallBase>(&I);
        if (CB == nullptr)
          continue;

        Function *Callee = CB->getCalledFunction();
        if (Callee && Callee->hasFnAttribute("uses-fp64-math")) {
          F->addFnAttr("uses-fp64-math");
          modified = true;

          markForRemoval(Callee);
          markedCallee = true;
        }
      }
    }

    if (markedCallee || F->hasFnAttribute("uses-fp64-math")) {
      markForRemoval(F);
    }
  }
  return modified;
}

static void poisonKernel(Function *Kernel) {
  Module *M = Kernel->getParent();
  LLVMContext &Ctx = M->getContext();

  Kernel->removeFnAttr("uses-fp64-math");
  Kernel->addFnAttr(PoisonFP64Kernels::attributeName);
  Kernel->deleteBody();

  BasicBlock *EntryBB = BasicBlock::Create(Ctx, "entry", Kernel);
  ReturnInst::Create(Ctx, EntryBB);
}

bool PoisonFP64Kernels::finalize(CodeGenContext *ctx, IGCMD::MetaDataUtils *pMdUtils) {
  CodeGenContext *CGC = ctx;
  IGCMD::MetaDataUtils *MDUtils = pMdUtils;

  bool modified = false;
  for (Function *F : reverse(fp64FunctionsOrder)) {
    std::string msg = "Removing unsupported FP64 function: '" + F->getName().str() + "'";
    CGC->EmitWarning(msg.c_str(), F);

    if (isEntryFunc(MDUtils, F)) {
      poisonKernel(F);
    } else {
      F->eraseFromParent();
    }
    modified = true;
  }

  return modified;
}

bool PoisonFP64Kernels::run(Module &M, CallGraph &CG, CodeGenContext *ctx, IGCMD::MetaDataUtils *pMdUtils) {
  reset();
  bool modified = false;
  // Walk the call-graph SCCs bottom-up, the same order the CallGraphSCCPass framework uses.
  for (auto I = scc_begin(&CG); !I.isAtEnd(); ++I) {
    const std::vector<CallGraphNode *> &SCCNodes = *I;
    modified |= runOnSCCNodes(SCCNodes, ctx);
  }
  modified |= finalize(ctx, pMdUtils);
  return modified;
}

#if LLVM_VERSION_MAJOR >= 16
PreservedAnalyses PoisonFP64KernelsNPM::run(Module &M, ModuleAnalysisManager &AM) {
  bool changed = m_impl.run(M, AM.getResult<CallGraphAnalysis>(M), AM.getResult<CodeGenContextAnalysis>(M).Ctx,
                            AM.getResult<MetaDataUtilsAnalysis>(M).MdUtils);
  return changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
}
#endif // LLVM_VERSION_MAJOR >= 16
