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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "Compiler/Optimizer/OpenCLPasses/SetFastMathFlags/SetFastMathFlags.hpp"
#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/InstIterator.h>
#include <llvm/IR/IntrinsicInst.h>
#include "common/LLVMWarningsPop.hpp"
#include "llvm/Demangle/Demangle.h"

using namespace llvm;
using namespace IGC;
using namespace IGC::IGCMD;

// Register pass to igc-opt
#define PASS_FLAG "igc-set-fast-math-flags"
#define PASS_DESCRIPTION "Set llvm fast math flags according to compiler options"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(SetFastMathFlagsLPM, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_END(SetFastMathFlagsLPM, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char SetFastMathFlagsLPM::ID = 0;

// Function where we intentionally skip fast math optimizations
static const std::array UnoptimizedFuncNames = {"__ocl_svml_cos", "__ocl_svml_sin", "__spirv_ocl_exp",
                                                "__spirv_ocl_native_exp", "__spirv_ocl_erf"};

SetFastMathFlagsLPM::SetFastMathFlagsLPM() : ModulePass(ID), m_impl() {
  initializeSetFastMathFlagsLPMPass(*PassRegistry::getPassRegistry());
}

SetFastMathFlagsLPM::SetFastMathFlagsLPM(FastMathFlags Mask) : ModulePass(ID), m_impl(Mask) {
  initializeSetFastMathFlagsLPMPass(*PassRegistry::getPassRegistry());
}

SetFastMathFlagsLPM::SetFastMathFlagsLPM(FastMathFlags Mask, bool skipUnsafeFpMathAttr)
    : ModulePass(ID), m_impl(Mask, skipUnsafeFpMathAttr) {
  initializeSetFastMathFlagsLPMPass(*PassRegistry::getPassRegistry());
}

bool SetFastMathFlags::run(Module &M, ModuleMetaData *pModMD) {
  auto hasFnAttributeSet = [](Function &F, StringRef Attr) {
    return F.hasFnAttribute(Attr) && F.getFnAttribute(Attr).getValueAsString() == "true";
  };

  const ModuleMetaData &modMD = *pModMD;
  bool changed = false;
  for (Function &F : M) {
    FastMathFlags fmfs;
    // We determine if we have to use the unsafe-fp-math attr or we can omit it
    bool unsafeFpMathAttrCheck = m_skipUnsafeFpMathAttr || hasFnAttributeSet(F, "unsafe-fp-math");
    // Fast relaxed math implies all other flags.
    if (modMD.compOpt.FastRelaxedMath && unsafeFpMathAttrCheck) {
      fmfs.setFast();
      fmfs &= m_Mask;
      changed |= setFlags(F, fmfs);
      continue;
    }
    // Unsafe math implies everything besides finite-math (ninf and nnan)
    if ((modMD.compOpt.UnsafeMathOptimizations)) {
      if (hasFnAttributeSet(F, "no-signed-zeros-fp-math"))
        fmfs.setNoSignedZeros();
      fmfs.setAllowContract();
      fmfs.setAllowReassoc();
      fmfs.setAllowReciprocal();
      fmfs.setApproxFunc();
    }
    // Finite math implies no infs and nans.
    if (modMD.compOpt.FiniteMathOnly) {
      if (hasFnAttributeSet(F, "no-infs-fp-math")) {
        fmfs.setNoInfs();
      }
      if (hasFnAttributeSet(F, "no-nans-fp-math")) {
        fmfs.setNoNaNs();
      }
    }
    if (modMD.compOpt.NoSignedZeros && hasFnAttributeSet(F, "no-signed-zeros-fp-math")) {
      fmfs.setNoSignedZeros();
    }
    fmfs &= m_Mask;
    changed |= setFlags(F, fmfs);
  }
  return changed;
}

static bool setMathFlags(IntrinsicInst *II) {
  switch (II->getIntrinsicID()) {
  case Intrinsic::pow:
  case Intrinsic::exp2:
  case Intrinsic::log:
  case Intrinsic::sqrt:
    II->setFast(true);
    return true;
  default:
    break;
  }
  return false;
}

bool SetFastMathFlags::setFlags(Function &F, FastMathFlags fmfs) {
  if (!fmfs.any()) {
    return false;
  }

  StringRef fName = F.getName();
  bool isUnoptimizedFunc = false;
  for (const auto &funcName : UnoptimizedFuncNames) {
    if (fName.contains(funcName)) {
      std::string MangledNameForDemangling = fName.str();
      std::string DemangledName = demangle(MangledNameForDemangling);

      size_t OpenParen = DemangledName.find('(');
      if (OpenParen != std::string::npos) {
        DemangledName = DemangledName.substr(0, OpenParen);
      }

      StringRef DemangledNameRef = DemangledName;
      if (DemangledNameRef == funcName) {
        isUnoptimizedFunc = true;
        break;
      }
    }
  }
  if (isUnoptimizedFunc)
    return false;

  bool changed = false;
  for (inst_iterator i = inst_begin(&F), e = inst_end(&F); i != e; ++i) {
    unsigned int op = i->getOpcode();
    switch (op) {
    case Instruction::FNeg:
    case Instruction::FAdd:
    case Instruction::FSub:
    case Instruction::FMul:
    case Instruction::FDiv:
    case Instruction::FRem:
      i->setFastMathFlags(fmfs); // this actually does an OR between flags.
      changed = true;
      break;
    case Instruction::Call:
      if (auto II = dyn_cast<IntrinsicInst>(&*i)) {
        if (fmfs.isFast())
          changed |= setMathFlags(II);
      } else if (isa<FPMathOperator>(&*i)) {
        i->setFastMathFlags(fmfs);
        changed = true;
      }
      break;
    default:
      break;
    }
  }
  return changed;
}

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