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

Copyright (C) 2017-2022 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 <llvmWrapper/IR/IRBuilder.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include "common/LLVMWarningsPop.hpp"
#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(ExtensionFuncsResolution, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_END(ExtensionFuncsResolution, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char ExtensionFuncsResolution::ID = 0;

ExtensionFuncsResolution::ExtensionFuncsResolution() : FunctionPass(ID), m_implicitArgs()
{
    initializeExtensionFuncsResolutionPass(*PassRegistry::getPassRegistry());
}

bool ExtensionFuncsResolution::runOnFunction(Function& F)
{
    m_changed = false;
    m_implicitArgs = ImplicitArgs(F, getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils());
    visit(F);
    return m_changed;
}

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

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

    ImplicitArg::ArgType argType;
    if (funcName.equals(ExtensionFuncsAnalysis::VME_MB_BLOCK_TYPE))
    {
        argType = ImplicitArg::VME_MB_BLOCK_TYPE;
    }
    else if (funcName.equals(ExtensionFuncsAnalysis::VME_SUBPIXEL_MODE))
    {
        argType = ImplicitArg::VME_SUBPIXEL_MODE;
    }
    else if (funcName.equals(ExtensionFuncsAnalysis::VME_SAD_ADJUST_MODE))
    {
        argType = ImplicitArg::VME_SAD_ADJUST_MODE;
    }
    else if (funcName.equals(ExtensionFuncsAnalysis::VME_SEARCH_PATH_TYPE))
    {
        argType = ImplicitArg::VME_SEARCH_PATH_TYPE;
    }
    else if (funcName.startswith(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 = PointerType::get(retType, 0);
        auto bitcastInst = builder.CreateBitCast(CI.getArgOperand(0), ptrType);
        auto ret = builder.CreateLoad(bitcastInst);
        CI.replaceAllUsesWith(ret);
        CI.eraseFromParent();
        return;
    }
    else if (funcName.startswith(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;
}
