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

Copyright (C) 2017-2026 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "Compiler/Optimizer/OpenCLPasses/ExtensionFuncs/ExtensionFuncResolution.hpp"
#include "Compiler/Optimizer/OpenCLPasses/ExtensionFuncs/ExtensionFuncsAnalysis.hpp"
#include "Compiler/IGCPassSupport.h"
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include "common/LLVMWarningsPop.hpp"
#include <llvmWrapper/IR/IRBuilder.h>
#include "llvmWrapper/IR/DerivedTypes.h"
#include "Probe/Assertion.h"

using namespace llvm;
using namespace IGC;

// Register pass to igc-opt
#define PASS_FLAG "igc-extension-funcs-resolution"
#define PASS_DESCRIPTION "Resolves extension function"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(ExtensionFuncsResolutionLPM, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_END(ExtensionFuncsResolutionLPM, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char ExtensionFuncsResolutionLPM::ID = 0;

ExtensionFuncsResolutionLPM::ExtensionFuncsResolutionLPM() : FunctionPass(ID) {
  initializeExtensionFuncsResolutionLPMPass(*PassRegistry::getPassRegistry());
}

bool ExtensionFuncsResolution::runOnFunction(Function &F, IGCMD::MetaDataUtils *pMdUtils, IGC::ModuleMetaData *pModMD) {
  m_changed = false;
  m_implicitArgs = ImplicitArgs(F, pMdUtils, pModMD);
  visit(F);
  return m_changed;
}

#if LLVM_VERSION_MAJOR >= 16
PreservedAnalyses ExtensionFuncsResolutionNPM::run(Module &M, ModuleAnalysisManager &AM) {
  auto *pMdUtils = AM.getResult<MetaDataUtilsAnalysis>(M).MdUtils;
  auto *pModMD = AM.getResult<MetaDataUtilsAnalysis>(M).ModMD;
  ExtensionFuncsResolution impl;
  bool changed = false;
  for (Function &F : M) {
    if (F.isDeclaration())
      continue;
    changed |= impl.runOnFunction(F, pMdUtils, pModMD);
  }
  return changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
}
#endif // LLVM_VERSION_MAJOR >= 16

void ExtensionFuncsResolution::visitCallInst(CallInst &CI) {
  if (!CI.getCalledFunction()) {
    return;
  }

  StringRef funcName = CI.getCalledFunction()->getName();
  Function &F = *(CI.getParent()->getParent());

  ImplicitArg::ArgType argType;
  if (funcName == (ExtensionFuncsAnalysis::VME_MB_BLOCK_TYPE)) {
    argType = ImplicitArg::VME_MB_BLOCK_TYPE;
  } else if (funcName == (ExtensionFuncsAnalysis::VME_SUBPIXEL_MODE)) {
    argType = ImplicitArg::VME_SUBPIXEL_MODE;
  } else if (funcName == (ExtensionFuncsAnalysis::VME_SAD_ADJUST_MODE)) {
    argType = ImplicitArg::VME_SAD_ADJUST_MODE;
  } else if (funcName == (ExtensionFuncsAnalysis::VME_SEARCH_PATH_TYPE)) {
    argType = ImplicitArg::VME_SEARCH_PATH_TYPE;
  } else if (IGCLLVM::starts_with(funcName, ExtensionFuncsAnalysis::VME_HELPER_GET_HANDLE)) {
    // Load from the opaque vme pointer and return the a vector with values.
    IGC_ASSERT(IGCLLVM::getNumArgOperands(&CI) == 1);
    IGCLLVM::IRBuilder<> builder(&CI);
    Type *retType = CI.getType();
    IGC_ASSERT(retType->isVectorTy() || retType->isIntegerTy());
    PointerType *ptrType = IGCLLVM::PointerType::get(retType, 0);
    auto bitcastInst = builder.CreateBitCast(CI.getArgOperand(0), ptrType);
    auto ret = builder.CreateLoad(retType, bitcastInst);
    CI.replaceAllUsesWith(ret);
    CI.eraseFromParent();
    return;
  } else if (IGCLLVM::starts_with(funcName, ExtensionFuncsAnalysis::VME_HELPER_GET_AS)) {
    // Store the VME values and return an opaque vme pointer.
    IGC_ASSERT(IGCLLVM::getNumArgOperands(&CI) == 1);
    IGCLLVM::IRBuilder<> builder(&*CI.getParent()->getParent()->begin()->getFirstInsertionPt());
    Type *retType = CI.getType();
    Value *arg = CI.getArgOperand(0);
    IGC_ASSERT(arg->getType()->isVectorTy() || arg->getType()->isIntegerTy());
    AllocaInst *allocaInst = builder.CreateAlloca(arg->getType());
    builder.SetInsertPoint(&CI);
    builder.CreateStore(arg, allocaInst);
    Value *bitcastInst = builder.CreateBitCast(allocaInst, retType);
    CI.replaceAllUsesWith(bitcastInst);
    CI.eraseFromParent();
    return;
  } else {
    // Non VME function, do nothing
    return;
  }

  Value *vmeRes = m_implicitArgs.getImplicitArg(F, argType);

  // Replace original VME call instruction with the appropriate argument, example:

  // Recieves:
  // call i32 @__builtin_IB_vme_subpixel_mode()

  // Creates:
  // %vmeSubpixelMode

  CI.replaceAllUsesWith(vmeRes);
  CI.eraseFromParent();

  m_changed = true;
}
