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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "AdaptorCommon/ImplicitArgs.hpp"
#include "Compiler/Optimizer/OpenCLPasses/LocalBuffers/InlineLocalsResolution.hpp"
#include "Compiler/CodeGenPublic.h"
#include "Compiler/IGCPassSupport.h"
#include "Compiler/CISACodeGen/helper.h"
#include "Compiler/DebugInfo/Utils.h"
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Module.h>
#include <llvm/IR/Instructions.h>
#include <llvmWrapper/Support/Alignment.h>
#include "common/LLVMWarningsPop.hpp"
#include "Probe/Assertion.h"

#include <unordered_set>

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


// Register pass to igc-opt
#define PASS_FLAG "igc-resolve-inline-locals"
#define PASS_DESCRIPTION "Resolve inline local variables/buffers"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(InlineLocalsResolution, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_END(InlineLocalsResolution, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char InlineLocalsResolution::ID = 0;
const llvm::StringRef BUILTIN_MEMPOOL = "__builtin_IB_AllocLocalMemPool";

InlineLocalsResolution::InlineLocalsResolution() :
    ModulePass(ID), m_pGV(nullptr)
{
    initializeInlineLocalsResolutionPass(*PassRegistry::getPassRegistry());
}

const unsigned int InlineLocalsResolution::VALID_LOCAL_HIGH_BITS = 0x10000000;

static bool useAsPointerOnly(Value* V) {
    IGC_ASSERT_MESSAGE(V->getType()->isPointerTy(), "Expect the input value is a pointer!");

    SmallSet<PHINode*, 8> VisitedPHIs;
    SmallVector<Value*, 16> WorkList;
    WorkList.push_back(V);

    StoreInst* ST = nullptr;
    PHINode* PN = nullptr;
    while (!WorkList.empty()) {
        Value* Val = WorkList.pop_back_val();
        for (auto* U : Val->users()) {
            Operator* Op = dyn_cast<Operator>(U);
            if (!Op)
                continue;
            switch (Op->getOpcode()) {
            default:
                // Bail out for unknown operations.
                return false;
            case Instruction::Store:
                ST = cast<StoreInst>(U);
                // Bail out if it's used as the value operand.
                if (ST->getValueOperand() == Val)
                    return false;
                // FALL THROUGH
            case Instruction::Load:
                // Safe use in LD/ST as pointer only.
                continue;
            case Instruction::PHI:
                PN = cast<PHINode>(U);
                // Skip if it's already visited.
                if (!VisitedPHIs.insert(PN).second)
                    continue;
                // FALL THROUGH
            case Instruction::BitCast:
            case Instruction::Select:
            case Instruction::GetElementPtr:
                // Need to check their usage further.
                break;
            }
            WorkList.push_back(U);
        }
    }

    return true;
}

bool InlineLocalsResolution::runOnModule(Module& M)
{
    MetaDataUtils* pMdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
    ModuleMetaData* modMD = getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
    if (!modMD->compOpt.OptDisable)
      filterGlobals(M);
    // Compute the offset of each inline local in the kernel,
    // and their total size.
    llvm::MapVector<Function*, unsigned int> sizeMap;
    collectInfoOnSharedLocalMem(M);
    computeOffsetList(M, sizeMap);

    LLVMContext& C = M.getContext();

    for (Function& F : M)
    {
        if (F.isDeclaration() || !isEntryFunc(pMdUtils, &F))
        {
            continue;
        }

        unsigned int totalSize = 0;

        // Get the offset at which local arguments start
        auto sizeIter = sizeMap.find(&F);
        if (sizeIter != sizeMap.end())
        {
            totalSize += sizeIter->second;
        }

        // Set the high 16 bits to a non-0 value.
        totalSize = (totalSize & 0xFFFF);

        bool IsFirstSLMArgument = true;
        for (Function::arg_iterator A = F.arg_begin(), AE = F.arg_end(); A != AE; ++A)
        {
            Argument* arg = &(*A);
            PointerType* ptrType = dyn_cast<PointerType>(arg->getType());
            // Check that this is a pointer
            if (!ptrType)
            {
                continue;
            }

            // To the local address space
            if (ptrType->getAddressSpace() != ADDRESS_SPACE_LOCAL)
            {
                continue;
            }

            // Which is used
            if (arg->use_empty())
            {
                continue;
            }

            bool UseAsPointerOnly = useAsPointerOnly(arg);
            unsigned Offset = totalSize;
            if (!UseAsPointerOnly)
                Offset |= VALID_LOCAL_HIGH_BITS;

            if (IsFirstSLMArgument) {
                auto BufType = ArrayType::get(Type::getInt8Ty(M.getContext()), 0);
                auto ExtSLM = new GlobalVariable(M, BufType, false, GlobalVariable::ExternalLinkage, nullptr,
                    F.getName() + "-ExtSLM", nullptr, GlobalVariable::ThreadLocalMode::NotThreadLocal,
                    ADDRESS_SPACE_LOCAL);
                auto NewPtr = ConstantExpr::getBitCast(ExtSLM, arg->getType());
                arg->replaceAllUsesWith(NewPtr);
                // Update MD.
                LocalOffsetMD localOffset;
                localOffset.m_Var = ExtSLM;
                localOffset.m_Offset = Offset;
                modMD->FuncMD[&F].localOffsets.push_back(localOffset);

                IGC::appendToUsed(M, ExtSLM);
                IsFirstSLMArgument = false;
            }
            else {
                // FIXME: The following code should be removed as well by
                // populating similar adjustment in prolog during code
                // emission.
                // Ok, now we need to add an offset, in bytes, which is equal to totalSize.
                // Bitcast to i8*, GEP, bitcast back to original type.
                Value* sizeConstant = ConstantInt::get(Type::getInt32Ty(C), Offset);
                SmallVector<Value*, 1> idx(1, sizeConstant);
                Instruction* pInsertBefore = &(*F.begin()->getFirstInsertionPt());
                Type* pCharType = Type::getInt8Ty(C);
                Type* pLocalCharPtrType = pCharType->getPointerTo(ADDRESS_SPACE_LOCAL);
                Instruction* pCharPtr = BitCastInst::CreatePointerCast(arg, pLocalCharPtrType, "localToChar", pInsertBefore);
                Value* pMovedCharPtr = GetElementPtrInst::Create(pCharType, pCharPtr, idx, "movedLocal", pInsertBefore);

                Value* pMovedPtr = CastInst::CreatePointerCast(pMovedCharPtr, ptrType, "charToLocal", pInsertBefore);

                // Running over arg users and use replaceUsesOfWith to fix them is not enough,
                // because it does not cover the usage of arg in metadata (e.g. for debug info intrinsic).
                // Thus, replaceAllUsesWith should be used in order to fix also debug info.
                arg->replaceAllUsesWith(pMovedPtr);
                // The above operation changed also the "arg" operand in "charPtr" to "movedPtr"
                // Thus, we need to fix it back (otherwise the LLVM IR will be invalid)
                pCharPtr->replaceUsesOfWith(pMovedPtr, arg);
            }
        }
    }

    return true;
}

void InlineLocalsResolution::filterGlobals(Module& M)
{
    // This data structure saves all the unused nodes,
    // including the global variable definition itself, as well as all successive recursive user nodes,
    // in all the def-use trees corresponding to all the global variables in the entire Module.
    std::unordered_set<Value*> unusedNodes_forModule;

    // let's loop all global variables
    for (Module::global_iterator I = M.global_begin(), E = M.global_end(); I != E; ++I)
    {
        // We only care about global variables, not other globals.
        GlobalVariable* globalVar = dyn_cast<GlobalVariable>(&*I);
        if (!globalVar)
        {
            continue;
        }

        PointerType* ptrType = cast<PointerType>(globalVar->getType());
        // We only care about local address space here.
        if (ptrType->getAddressSpace() != ADDRESS_SPACE_LOCAL)
        {
            continue;
        }

        // If the globalVar is determined to be unused,
        // this data structure saves the globalVar,
        // as well as all successive recursive user nodes in that def-use tree.
        std::unordered_set<Value*> unusedNodes_forOne;
        if (unusedGlobal(globalVar, unusedNodes_forOne))
            unusedNodes_forModule.insert(unusedNodes_forOne.begin(), unusedNodes_forOne.end());
    }

    // We only remove all the unused nodes for this Module,
    // after we are done processing all the global variables for the entire Module,
    // to prevent iterators becoming invalidated when elements get removed from the ilist.
    for (auto& element : unusedNodes_forModule) {
        // for all unused Values,
        //   replace all uses with undefs
        //   delete the values
        if (Instruction* node = dyn_cast<Instruction>(element)) {
            Type* Ty = node->getType();
            if (!Ty->isVoidTy())
                node->replaceAllUsesWith(UndefValue::get(Ty));
            node->eraseFromParent();
        }
        else if (GlobalVariable* node = dyn_cast<GlobalVariable>(element)) {
            Type* Ty = node->getType();
            if (!Ty->isVoidTy())
                node->replaceAllUsesWith(UndefValue::get(Ty));
            node->eraseFromParent();
        }
        // All other types of nodes are ignored.
    }
}

bool InlineLocalsResolution::unusedGlobal(Value* V, std::unordered_set<Value*>& unusedNodes)
{
    for (Value::user_iterator U = V->user_begin(), UE = V->user_end(); U != UE; ++U)
    {
        if (GlobalVariable* globalVar = dyn_cast<GlobalVariable>(*U)) {
            if (!unusedGlobal(*U, unusedNodes))
                return false;
        }
        else if (GetElementPtrInst* gep = dyn_cast<GetElementPtrInst>(*U)) {
            if (!unusedGlobal(*U, unusedNodes))
                return false;
        }
        else if (BitCastInst* bitcast = dyn_cast<BitCastInst>(*U)) {
            if (!unusedGlobal(*U, unusedNodes))
                return false;
        }
        else if (StoreInst* store = dyn_cast<StoreInst>(*U)) {
            if (store->isUnordered()) {
                if (store->getPointerOperand() == V) {
                    if (!unusedGlobal(*U, unusedNodes))
                        return false;
                }
                else if (store->getValueOperand() == V) {
                    return false;
                }
            }
            else {
                return false;
            }
        }
        else {  // some other instruction
            return false;
        }
    }
    // add an unused node to the data structure
    unusedNodes.insert(V);
    return true;
}

void InlineLocalsResolution::collectInfoOnSharedLocalMem(Module& M)
{
    const auto pCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
    // first we collect SLM usage on GET_MEMPOOL_PTR
    if (M.getFunction(BUILTIN_MEMPOOL) != nullptr)
    {
        const GT_SYSTEM_INFO platform = pCtx->platform.GetGTSystemInfo();

        SmallVector<CallInst*, 8> callsToReplace;
        unsigned maxBytesOnModule = 0;
        unsigned maxAlignOnModule = 0;

        unsigned int maxWorkGroupSize = 448;
        maxWorkGroupSize = pCtx->platform.getOfflineCompilerMaxWorkGroupSize();
        if (pCtx->platform.enableMaxWorkGroupSizeCalculation() &&
            platform.EUCount != 0 && platform.SubSliceCount != 0)
        {
            unsigned int maxNumEUsPerSubSlice = platform.EuCountPerPoolMin;
            if (platform.EuCountPerPoolMin == 0 || pCtx->platform.supportPooledEU())
            {
                maxNumEUsPerSubSlice = platform.EUCount / platform.SubSliceCount;
            }
            const unsigned int numThreadsPerEU = platform.ThreadCount / platform.EUCount;
            unsigned int simdSizeUsed = 8;
            unsigned int maxWS = maxNumEUsPerSubSlice * numThreadsPerEU * simdSizeUsed;
            if (!iSTD::IsPowerOfTwo(maxWS))
            {
                maxWS = iSTD::RoundPower2((DWORD)maxWS) >> 1;
            }
            maxWorkGroupSize = std::min(maxWS, 1024u);
        }

        // scan inst to collect all call instructions

        for (Function& F : M)
        {
            if (F.isDeclaration())
            {
                continue;
            }

            unsigned maxBytesOnFunc = 0;
            for (auto I = inst_begin(&F), IE = inst_end(&F); I != IE; ++I)
            {
                Instruction* inst = &(*I);
                if (CallInst * CI = dyn_cast<CallInst>(inst))
                {
                    Function* pFunc = CI->getCalledFunction();
                    if (pFunc && pFunc->getName().equals(BUILTIN_MEMPOOL))
                    {
                        // should always be called with constant operands
                        IGC_ASSERT(isa<ConstantInt>(CI->getArgOperand(0)));
                        IGC_ASSERT(isa<ConstantInt>(CI->getArgOperand(1)));
                        IGC_ASSERT(isa<ConstantInt>(CI->getArgOperand(2)));

                        const unsigned int allocAllWorkgroups = unsigned(cast<ConstantInt>(CI->getArgOperand(0))->getZExtValue());
                        const unsigned int numAdditionalElements = unsigned(cast<ConstantInt>(CI->getArgOperand(1))->getZExtValue());
                        const unsigned int elementSize = unsigned(cast<ConstantInt>(CI->getArgOperand(2))->getZExtValue());

                        unsigned int numElements = numAdditionalElements;
                        if (allocAllWorkgroups)
                        {
                            numElements += maxWorkGroupSize;
                        }
                        const unsigned int size = numElements * elementSize;
                        const unsigned int align = elementSize;

                        maxBytesOnFunc = std::max(maxBytesOnFunc, size);
                        maxBytesOnModule = std::max(maxBytesOnModule, size);
                        maxAlignOnModule = std::max(maxAlignOnModule, align);

                        callsToReplace.push_back(CI);
                    }
                }
            }
            if (maxBytesOnFunc != 0)
            {
                m_FuncToMemPoolSizeMap[&F] = maxBytesOnFunc;
            }
        }

        if (!callsToReplace.empty())
        {

            Type* bufType = ArrayType::get(Type::getInt8Ty(M.getContext()), uint64_t(maxBytesOnModule));

            m_pGV = new GlobalVariable(M, bufType, false,
                GlobalVariable::ExternalLinkage, ConstantAggregateZero::get(bufType),
                "GenSLM.LocalMemPoolOnGetMemPoolPtr",
                nullptr,
                GlobalVariable::ThreadLocalMode::NotThreadLocal,
                ADDRESS_SPACE_LOCAL);

            m_pGV->setAlignment(IGCLLVM::getCorrectAlign(maxAlignOnModule));

            for (auto call : callsToReplace)
            {
                CastInst* cast =
                    new BitCastInst(
                        m_pGV,
                        call->getCalledFunction()->getReturnType(),
                        "mempoolcast",
                        call);

                cast->setDebugLoc(call->getDebugLoc());

                call->replaceAllUsesWith(cast);
                call->eraseFromParent();
            }
        }
    }

    // let's loop all global variables
    for (Module::global_iterator I = M.global_begin(), E = M.global_end(); I != E; ++I)
    {
        // We only care about global variables, not other globals.
        GlobalVariable* globalVar = dyn_cast<GlobalVariable>(&*I);
        if (!globalVar)
        {
            continue;
        }

        PointerType* ptrType = dyn_cast<PointerType>(globalVar->getType());
        IGC_ASSERT_MESSAGE(ptrType, "The type of a global variable must be a pointer type");
        if (!ptrType)
        {
            continue;
        }

        // We only care about local address space here.
        if (ptrType->getAddressSpace() != ADDRESS_SPACE_LOCAL)
        {
            continue;
        }

        // For each SLM buffer, set section to avoid alignment changing by llvm.
        // Add external linkage and DSO scope information.
        globalVar->setLinkage(GlobalValue::ExternalLinkage);
        globalVar->setDSOLocal(false);
        globalVar->setSection("localSLM");

        // Find the functions which this globalVar belongs to....
        for (Value::user_iterator U = globalVar->user_begin(), UE = globalVar->user_end(); U != UE; ++U)
        {
            Instruction* user = dyn_cast<Instruction>(*U);
            if (!user)
            {
                continue;
            }

            Function* parentF = user->getParent()->getParent();
            bool emitError = false;
            if (pCtx->type == ShaderType::OPENCL_SHADER)
            {
                // If this option is passed, emit error when extern functions use local SLM
                auto ClContext = static_cast<OpenCLProgramContext*>(pCtx);
                emitError = ClContext->m_Options.EmitErrorsForLibCompilation;
            }
            if (parentF->hasFnAttribute("referenced-indirectly") && emitError)
            {
                IGC_ASSERT_MESSAGE(0, "Cannot reference localSLM in indirectly-called functions");
                getAnalysis<CodeGenContextWrapper>().getCodeGenContext()->EmitError("Cannot reference localSLM in indirectly-called functions", globalVar);
                return;
            }
            m_FuncToVarsMap[parentF].insert(globalVar);
        }
    }

    // set debugging info, and insert mov inst.
    for (const auto& I : m_FuncToVarsMap)
    {
        Function* userFunc = I.first;
        for (auto* G : I.second)
        {
            Instruction * pInsertBefore = &(*userFunc->begin()->getFirstInsertionPt());
            TODO("Should inline local buffer points to origin offset 'globalVar' or to fixed offset 'pMovedPtr'?");
            Utils::UpdateGlobalVarDebugInfo(G, G, pInsertBefore, true);
        }
    }
}

void InlineLocalsResolution::computeOffsetList(Module& M, llvm::MapVector<Function*, unsigned int>& sizeMap)
{
    llvm::MapVector<Function*, llvm::MapVector<GlobalVariable*, unsigned int>> offsetMap;
    MetaDataUtils* pMdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
    ModuleMetaData* modMD = getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
    DataLayout DL = M.getDataLayout();
    CallGraph& CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();

    if (m_FuncToVarsMap.empty())
    {
        return;
    }

    // let's traverse the CallGraph to calculate the local
    // variables of kernel from all user functions.
    m_chkSet.clear();
    for (auto& N : CG)
    {
        Function* f = N.second->getFunction();
        if (!f || f->isDeclaration() || m_chkSet.find(f) != m_chkSet.end()) continue;
        traverseCGN(*N.second);
    }

    // set up the offsetMap;
    for (const auto& I : m_FuncToVarsMap)
    {
        Function* F = I.first;

        // loop through all global variables
        for (auto G : I.second)
        {
            auto itr = sizeMap.find(F);
            unsigned int offset = itr == sizeMap.end() ? 0 : itr->second;
#if LLVM_VERSION_MAJOR < 11
            offset = iSTD::Align(offset, DL.getPreferredAlignment(G));
#else
            offset = iSTD::Align(offset, (unsigned)DL.getPreferredAlign(G).value());
#endif
            // Save the offset of the current local
            // (set the high bits to be non-0 here too)
            offsetMap[F][G] = (offset & 0xFFFF);

            // And the total size after this local is added
            PointerType* ptrType = dyn_cast<PointerType>(G->getType());
            Type* varType = ptrType->getPointerElementType();
            if (G == m_pGV)
            {
                // it is GetMemPoolPtr usage
                offset += m_FuncToMemPoolSizeMap[F];
            }
            else
            {
                offset += (unsigned int)DL.getTypeAllocSize(varType);
            }
            sizeMap[F] = offset;
        }
    }

    // Ok, we've collected the information, now write it into the MD.
    for (auto& iter : sizeMap)
    {
        // ignore non-entry functions.
        if (!isEntryFunc(pMdUtils, iter.first))
        {
            continue;
        }

        // If this function doesn't have any locals, no need for MD.
        if (iter.second == 0)
        {
            continue;
        }

        // We need the total size to have at least 32-byte alignment.
        // This is because right after the space allocated to the inline locals,
        // we are going to have inline parameters. So, we need to make sure the
        // first local parameter is appropriately aligned, which, at worst,
        // can be 256 bits.
        iter.second = iSTD::Align(iter.second, 32);

        // Add the size information of this function
        modMD->FuncMD[iter.first].localSize = iter.second;

        // And now the offsets.
        for (const auto& offsetIter : offsetMap[iter.first])
        {
            unsigned Offset = offsetIter.second;
            if (!useAsPointerOnly(offsetIter.first))
                Offset |= VALID_LOCAL_HIGH_BITS;

            LocalOffsetMD localOffset;
            localOffset.m_Var = offsetIter.first;
            localOffset.m_Offset = Offset;
            modMD->FuncMD[iter.first].localOffsets.push_back(localOffset);
        }
    }
    pMdUtils->save(M.getContext());
}

void InlineLocalsResolution::traverseCGN(const llvm::CallGraphNode& CGN)
{
    Function* f = CGN.getFunction();

    // mark this function
    m_chkSet.insert(f);

    for (const auto& N : CGN)
    {
        Function* sub = N.second->getFunction();
        if (!sub || sub->isDeclaration()) continue;

        // we reach here, because there is sub-function inside the node
        if (m_chkSet.find(sub) == m_chkSet.end())
        {
            // this sub-routine is not visited before.
            // visit it first
            traverseCGN(*N.second);
        }

        // the sub-routine was visited before, collect information

        // count each global on this sub-routine
        GlobalVariableSet& GS_f = m_FuncToVarsMap[f];
        const GlobalVariableSet& GS_sub = m_FuncToVarsMap[sub];
        GS_f.insert(GS_sub.begin(), GS_sub.end());

        // automatic storages
        if (m_FuncToMemPoolSizeMap.find(sub) != m_FuncToMemPoolSizeMap.end())
        {
            // this sub-function has automatic storage
            if (m_FuncToMemPoolSizeMap.find(f) != m_FuncToMemPoolSizeMap.end())
            {
                // caller has its own memory pool size, choose the max
                m_FuncToMemPoolSizeMap[f] = std::max(m_FuncToMemPoolSizeMap[f], m_FuncToMemPoolSizeMap[sub]);
            }
            else
            {
                m_FuncToMemPoolSizeMap[f] = m_FuncToMemPoolSizeMap[sub];
            }
        }
    }
}