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

Copyright (C) 2022-2023 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "AdaptorCommon/ImplicitArgs.hpp"
#include "Compiler/Optimizer/OpenCLPasses/ScalarArgAsPointer/ScalarArgAsPointer.hpp"
#include "Compiler/CISACodeGen/OpenCLKernelCodeGen.hpp"
#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Module.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/GetElementPtrTypeIterator.h>
#include "llvm/Support/Debug.h"
#include "common/LLVMWarningsPop.hpp"

#include "Probe/Assertion.h"

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

#define DEBUG_TYPE "igc-scalar-arg-as-pointer-analysis"

// Register pass to igc-opt
#define PASS_FLAG "igc-scalar-arg-as-pointer-analysis"
#define PASS_DESCRIPTION "Analyzes scalar kernel arguments used for global memory access"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(ScalarArgAsPointerAnalysis, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_END(ScalarArgAsPointerAnalysis, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char ScalarArgAsPointerAnalysis::ID = 0;

ScalarArgAsPointerAnalysis::ScalarArgAsPointerAnalysis() : ModulePass(ID)
{
    initializeScalarArgAsPointerAnalysisPass(*PassRegistry::getPassRegistry());
}

bool ScalarArgAsPointerAnalysis::runOnModule(Module& M)
{
    DL = &M.getDataLayout();

    MDU = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();

    bool changed = false;

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

        if (!isEntryFunc(MDU, &F))
            continue;

        changed |= analyzeFunction(F);
    }

    // Update LLVM metadata based on IGC MetadataUtils
    if (changed)
        MDU->save(M.getContext());

    return changed;
}

bool ScalarArgAsPointerAnalysis::analyzeFunction(llvm::Function& F)
{
    m_matchingArgs.clear();
    m_visitedInst.clear();
    m_allocas.clear();
    m_currentFunction = &F;

    LLVM_DEBUG(
        dbgs() << "running for function " << F.getName() << "\n");

    visit(F);

    if (m_matchingArgs.empty())
        return false;

    FunctionMetaData& funcMD = getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData()->FuncMD[&F];

    for (auto it = m_matchingArgs.begin(); it != m_matchingArgs.end(); ++it)
        funcMD.m_OpenCLArgScalarAsPointers.insert((*it)->getArgNo());

    return true;
}

void ScalarArgAsPointerAnalysis::visitStoreInst(llvm::StoreInst& I)
{
    analyzeStoredArg(I);
    analyzePointer(I.getPointerOperand());
}

void ScalarArgAsPointerAnalysis::visitLoadInst(llvm::LoadInst& I)
{
    analyzePointer(I.getPointerOperand());
}

void ScalarArgAsPointerAnalysis::visitCallInst(CallInst& CI)
{
    if (auto I = dyn_cast<GenIntrinsicInst>(&CI))
        return visitGenIntrinsic(*I);

    // If function takes pointer as argument, assume it is used for access.
    for (auto it = CI.op_begin(); it != CI.op_end(); ++it)
    {
        if (isa< PointerType>((*it)->getType()))
            analyzePointer(*it);
    }
}

void ScalarArgAsPointerAnalysis::visitGenIntrinsic(GenIntrinsicInst& I)
{
    GenISAIntrinsic::ID const id = I.getIntrinsicID();

    if (id == GenISAIntrinsic::GenISA_LSC2DBlockRead ||
        id == GenISAIntrinsic::GenISA_LSC2DBlockPrefetch ||
        id == GenISAIntrinsic::GenISA_LSC2DBlockWrite)
    {
        return analyzeValue(I.getOperand(0));
    }

    if (IsStatelessMemLoadIntrinsic(id) ||
        IsStatelessMemStoreIntrinsic(id) ||
        IsStatelessMemAtomicIntrinsic(I, id))
    {
        Value* V = GetBufferOperand(&I);

        if (!V || !isa<PointerType>(V->getType()))
            return;

        return analyzePointer(V);
    }
}

void ScalarArgAsPointerAnalysis::analyzePointer(llvm::Value* V)
{
    auto* type = dyn_cast<PointerType>(V->getType());

    IGC_ASSERT_MESSAGE(type, "Value should be a pointer");

    if (type->getAddressSpace() != ADDRESS_SPACE_GLOBAL && type->getAddressSpace() != ADDRESS_SPACE_GENERIC)
        return;

    analyzeValue(V);
}

void ScalarArgAsPointerAnalysis::analyzeValue(llvm::Value* V)
{
    std::shared_ptr<ScalarArgAsPointerAnalysis::ArgSet> args;

    if (auto* I = dyn_cast<Instruction>(V))
    {
        args = findArgs(I);
    }
    else
    {
        args = analyzeOperand(V);
    }

    if (args)
    {
        LLVM_DEBUG(
            for (auto a : *args)
            {
                dbgs() << "  access from pointer ";
                V->printAsOperand(dbgs(), false);
                dbgs() << " tracks to argument ";
                a->printAsOperand(dbgs(), false);
                dbgs() << "\n";
            });
        m_matchingArgs.insert(args->begin(), args->end());
    }
}

const std::shared_ptr<ScalarArgAsPointerAnalysis::ArgSet>
ScalarArgAsPointerAnalysis::findArgs(llvm::Instruction* inst)
{
    // Skip already visited instruction
    if (m_visitedInst.count(inst))
        return m_visitedInst[inst];

    // Mark as visited
    m_visitedInst.try_emplace(inst, nullptr);

    // Assume intrinsic are safe simple arithmetics.
    if (isa<CallInst>(inst) && !isa<GenIntrinsicInst>(inst))
        return nullptr;

    auto result = std::make_shared<ScalarArgAsPointerAnalysis::ArgSet>();

    if (LoadInst* LI = dyn_cast<LoadInst>(inst))
    {
        if (!findStoredArgs(*LI, *result))
            return nullptr; // (1) Found indirect access, fail search
    }
    else
    {
        // Iterate and trace back operands.
        auto begin = inst->operands().begin();
        auto end = inst->operands().end();

        if (isa<SelectInst>(inst))
        {
            // For select, skip condition operand (first arg)
            begin++;
        }
        else if (isa<GetElementPtrInst>(inst))
        {
            // For GEP, use only base pointer operand (first arg)
            end = begin + 1;
        }

        for (auto it = begin; it != end; ++it)
        {
            auto args = analyzeOperand(*it);

            if (args)
                result->insert(args->begin(), args->end());
        }

        if (result->empty())
            return nullptr; // propagate fail
    }

    m_visitedInst[inst] = result;
    return result;
}

const std::shared_ptr<ScalarArgAsPointerAnalysis::ArgSet>
ScalarArgAsPointerAnalysis::analyzeOperand(llvm::Value* op)
{
    auto result = std::make_shared<ScalarArgAsPointerAnalysis::ArgSet>();

    if (Argument* arg = dyn_cast<Argument>(op))
    {
        // Consider only integer arguments
        if (arg->getType()->getScalarType()->isIntegerTy())
        {
            result->insert(arg);
        }
        else
        {
            // (2) Found non-compatible argument, fail
            return nullptr;
        }
    }
    else if (Instruction* opInst = dyn_cast<Instruction>(op))
    {
        auto args = findArgs(opInst);

        if (!args)
            return nullptr; // propagate fail

        result->insert(args->begin(), args->end());
    }
    else if (GlobalValue* global = dyn_cast<GlobalValue>(op))
    {
        Argument* arg = analyzeGlobal(global);

        if (!arg)
            return nullptr; // propagate fail

        result->insert(arg);
    }

    return result;
}

llvm::Argument* ScalarArgAsPointerAnalysis::analyzeGlobal(llvm::GlobalValue* V)
{
    PointerType* type = dyn_cast<PointerType>(V->getType());
    if (!type)
        return nullptr;

    if (type->getAddressSpace() != ADDRESS_SPACE_GLOBAL)
        return nullptr;

    ImplicitArgs implicitArgs(*m_currentFunction, MDU);

    if (!implicitArgs.isImplicitArgExist(ImplicitArg::GLOBAL_BASE))
        return nullptr;

    if (m_currentFunction->arg_size() < implicitArgs.size())
        return nullptr;

    unsigned implicitArgsBaseIndex = m_currentFunction->arg_size() - implicitArgs.size();
    unsigned implicitArgIndex = implicitArgs.getArgIndex(ImplicitArg::GLOBAL_BASE);
    return std::next(m_currentFunction->arg_begin(), implicitArgsBaseIndex + implicitArgIndex);
}

void ScalarArgAsPointerAnalysis::analyzeStoredArg(llvm::StoreInst& SI)
{
    // Only track stores of kernel arguments.
    Argument* A = dyn_cast<Argument>(SI.getValueOperand());
    if (!A)
        return;

    AllocaInst* AI = nullptr;
    GetElementPtrInst* GEPI = nullptr;
    if (!findAllocaWithOffset(SI.getPointerOperand(), AI, GEPI))
        return;

    uint64_t totalOffset = 0;

    if (GEPI)
    {
        // For store instruction offset must be constant.
        APInt offset(DL->getIndexTypeSizeInBits(GEPI->getType()), 0);
        if (!GEPI->accumulateConstantOffset(*DL, offset) || offset.isNegative())
            return;
        totalOffset += offset.getZExtValue();
    }

    m_allocas[std::pair<llvm::AllocaInst*, uint64_t>(AI, totalOffset)] = A;
}

bool ScalarArgAsPointerAnalysis::findStoredArgs(llvm::LoadInst& LI, ArgSet& args)
{
    AllocaInst* AI = nullptr;
    GetElementPtrInst* GEPI = nullptr;
    if (!findAllocaWithOffset(LI.getPointerOperand(), AI, GEPI))
        return false;

    // It is possible one or more GEP operand is a variable index to array type.
    // In this case search for all possible offsets to alloca.
    using Offsets = SmallVector<uint64_t, 4>;
    Offsets offsets;
    offsets.push_back(0);

    if (GEPI)
    {
        for (gep_type_iterator GTI = gep_type_begin(GEPI), prevGTI = gep_type_end(GEPI); GTI != gep_type_end(GEPI); prevGTI = GTI++)
        {
            if (ConstantInt* C = dyn_cast<ConstantInt>(GTI.getOperand()))
            {
                if (C->isZero())
                    continue;

                uint64_t offset = 0;

                if (StructType* STy = GTI.getStructTypeOrNull())
                    offset = DL->getStructLayout(STy)->getElementOffset(int_cast<unsigned>(C->getZExtValue()));
                else
                    offset = C->getZExtValue() * DL->getTypeAllocSize(GTI.getIndexedType()); // array or vector

                for (auto it = offsets.begin(); it != offsets.end(); ++it)
                    *it += offset;
            }
            else
            {
                if (prevGTI == gep_type_end(GEPI))
                    return false; // variable index at first operand, should not happen

                // gep_type_iterator is used to query indexed type. For arrays this is type
                // of single element. To get array size, we need to do query for it at
                // previous iterator step (before stepping into type indexed by array).
                ArrayType* ATy = dyn_cast<ArrayType>(prevGTI.getIndexedType());
                if (!ATy)
                    return false;

                uint64_t arrayElements = ATy->getNumElements();

                uint64_t byteSize = DL->getTypeAllocSize(GTI.getIndexedType());

                Offsets tmp;
                for (auto i = 0; i < arrayElements; ++i)
                    for (auto it = offsets.begin(); it != offsets.end(); ++it)
                        tmp.push_back(*it + i * byteSize);

                offsets = std::move(tmp);
            }
        }
    }

    for (auto it = offsets.begin(); it != offsets.end(); ++it)
    {
        std::pair<llvm::AllocaInst*, uint64_t> key(AI, *it);
        if (m_allocas.count(key))
            args.insert(m_allocas[key]);
    }

    return !args.empty();
}

bool ScalarArgAsPointerAnalysis::findAllocaWithOffset(llvm::Value* V, llvm::AllocaInst*& outAI, llvm::GetElementPtrInst*& outGEPI)
{
    IGC_ASSERT_MESSAGE(dyn_cast<PointerType>(V->getType()), "Value should be a pointer");

    outGEPI = nullptr;
    Value* tmp = V;

    while (true)
    {
        if (BitCastInst* BCI = dyn_cast<BitCastInst>(tmp))
        {
            tmp = BCI->getOperand(0);
        }
        else if (GetElementPtrInst* GEPI = dyn_cast<GetElementPtrInst>(tmp))
        {
            if (outGEPI)
                return false; // only one GEP instruction is supported
            outGEPI = GEPI;
            tmp = GEPI->getPointerOperand();
        }
        else if (AllocaInst* AI = dyn_cast<AllocaInst>(tmp))
        {
            outAI = AI;
            return true;
        }
        else
        {
            return false;
        }
    }
}
