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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "common/debug/Debug.hpp"
#include "common/debug/Dump.hpp"
#include "Compiler/CISACodeGen/helper.h"
#include "Compiler/FixResourcePtr.hpp"
#include "Compiler/CodeGenPublic.h"
#include "Compiler/IGCPassSupport.h"
#include "GenISAIntrinsics/GenIntrinsics.h"
#include "GenISAIntrinsics/GenIntrinsicInst.h"
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/InstIterator.h>
#include "common/LLVMWarningsPop.hpp"
#include "Probe/Assertion.h"

using namespace llvm;
using namespace IGC;

// Register pass to igc-opt
#define PASS_FLAG "igc-fix-resource-ptr"
#define PASS_DESCRIPTION "Fix the usage of GetBufferPtr, no combination of GetBufferPtr and GetResourcePtr"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(FixResourcePtr, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_END(FixResourcePtr, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char FixResourcePtr::ID = 0;

FixResourcePtr::FixResourcePtr() : FunctionPass(ID)
{
    initializeFixResourcePtrPass(*PassRegistry::getPassRegistry());
}

bool FixResourcePtr::runOnFunction(llvm::Function& F)
{
    llvm::IRBuilder<> __builder(F.getContext());
    builder = &__builder;
    DL = &F.getParent()->getDataLayout();
    m_changed = false;
    curFunc = &F;

    std::vector<Instruction*> fixlist;
    // initialize worklist with all the GetBufferPtrs with immed resource/sampler id
    inst_iterator it = inst_begin(&F);
    inst_iterator  e = inst_end(&F);
    for (; it != e; ++it)
    {
        GenIntrinsicInst* inst = dyn_cast<GenIntrinsicInst>(&*it);
        if (inst && inst->getIntrinsicID() == GenISAIntrinsic::GenISA_GetBufferPtr)
        {
            unsigned int as = inst->getType()->getPointerAddressSpace();
            Value* bufferIndex = ResolveBufferIndex(inst->getOperand(0));
            if (isa<ConstantInt>(bufferIndex) && IGC::IsDirectIdx(as))
            {
                RemoveGetBufferPtr(inst, bufferIndex);
            }
            else
            {
                fixlist.push_back(inst);
            }
        }
    }

    // fix the load/store that uses pointer coming from a GetBufferPtr/GetElementPtr,
    // change it to intrinsic that directly uses buf-pointer
    while (!fixlist.empty())
    {
        Instruction* inst = fixlist.back();
        fixlist.pop_back();

        FindGetElementPtr(inst, inst);
    }
    while (!eraseList.empty())
    {
        Instruction* inst = eraseList.back();
        eraseList.pop_back();
        inst->eraseFromParent();
    }
    return m_changed;
}

void FixResourcePtr::RemoveGetBufferPtr(GenIntrinsicInst* bufPtr, Value* bufIdx)
{
    uint outAS = bufPtr->getType()->getPointerAddressSpace();
    uint origAS = outAS;
    BufferType bufType = (BufferType)(cast<ConstantInt>(bufPtr->getOperand(1))->getZExtValue());
    uint encodeAS = EncodeAS4GFXResource(*bufIdx, bufType);
    if (outAS != encodeAS &&
        (bufType == CONSTANT_BUFFER || bufType == RESOURCE || bufType == UAV))
    {
        // happens to OGL, need to fix if address-space encoding is wrong
        outAS = encodeAS;
    }

    std::vector<Instruction*> foldlist;
    foldlist.push_back(bufPtr);
    // fold instructions on the worklist to constant null-pointer value
    while (!foldlist.empty())
    {
        Instruction* inst = foldlist.back();
        foldlist.pop_back();

        PointerType* const instType = dyn_cast<PointerType>(inst->getType());
        IGC_ASSERT(nullptr != instType);
        Type* eltType = instType->getPointerElementType();
        PointerType* ptrType = PointerType::get(eltType, outAS);
        inst->mutateType(ptrType);
        // iterate all the uses, put bitcast on the worklist
        for (auto UI = inst->user_begin(), E = inst->user_end(); UI != E; ++UI)
        {
            if (BitCastInst * use = dyn_cast<BitCastInst>(*UI))
            {
                foldlist.push_back(use);
            }
            else if (GetElementPtrInst * gep = dyn_cast<GetElementPtrInst>(*UI))
            {
                Value* byteOffset = GetByteOffset(gep);
                builder->SetInsertPoint(gep);
                Value* int2ptr = builder->CreateIntToPtr(byteOffset, ptrType);
                gep->mutateType(ptrType);
                gep->replaceAllUsesWith(int2ptr);
                if (outAS != origAS)
                {
                    FixAddressSpaceInAllUses(int2ptr, outAS, origAS);
                }
            }
        }
        ConstantPointerNull* basePtr = ConstantPointerNull::get(ptrType);
        inst->replaceAllUsesWith(basePtr);
    }
}

void FixResourcePtr::FindGetElementPtr(Instruction* bufPtr, Instruction* searchPtr)
{
    // iterate all the uses, recursively find the GetElementPtr
    for (auto UI = searchPtr->user_begin(), E = searchPtr->user_end(); UI != E; ++UI)
    {
        if (BitCastInst * use = dyn_cast<BitCastInst>(*UI))
        {
            FindGetElementPtr(bufPtr, use);
        }
        else if (GetElementPtrInst * use = dyn_cast<GetElementPtrInst>(*UI))
        {
            FindLoadStore(bufPtr, use, use);
        }
        else if (LoadInst * ld = dyn_cast<LoadInst>(*UI))
        {
            // fix load
            Value* offsetValue = builder->getInt32(0);
            Value* loadIndexed = CreateLoadIntrinsic(ld, bufPtr, offsetValue);
            ld->replaceAllUsesWith(loadIndexed);
            eraseList.push_back(ld);
        }
        else if (StoreInst * st = dyn_cast<StoreInst>(*UI))
        {
            // fix store
            Value* offsetValue = builder->getInt32(0);
            Value* storeIndexed = CreateStoreIntrinsic(st, bufPtr, offsetValue);
            st->replaceAllUsesWith(storeIndexed);
            eraseList.push_back(st);
        }
    }
}

void FixResourcePtr::FindLoadStore(Instruction* bufPtr, Instruction* eltPtr, Instruction* searchPtr)
{
    // iterate all the uses, put bitcast on the worklist, recursively find the load/store
    for (auto UI = searchPtr->user_begin(), E = searchPtr->user_end(); UI != E; ++UI)
    {
        if (BitCastInst * use = dyn_cast<BitCastInst>(*UI))
        {
            FindLoadStore(bufPtr, eltPtr, use);
        }
        else if (LoadInst * use = dyn_cast<LoadInst>(*UI))
        {
            // fix load
            Value* offsetValue = GetByteOffset(eltPtr);
            Value* loadIndexed = CreateLoadIntrinsic(use, bufPtr, offsetValue);
            use->replaceAllUsesWith(loadIndexed);
            eraseList.push_back(use);
        }
        else if (StoreInst * use = dyn_cast<StoreInst>(*UI))
        {
            // fix store
            Value* offsetValue = GetByteOffset(eltPtr);
            Value* storeIndexed = CreateStoreIntrinsic(use, bufPtr, offsetValue);
            use->replaceAllUsesWith(storeIndexed);
            eraseList.push_back(use);
        }
    }
}

Value* FixResourcePtr::GetByteOffset(Instruction* eltPtr)
{
    IGC_ASSERT(eltPtr->getNumOperands() == 2);
    Value* ptrOp = eltPtr->getOperand(0);
    PointerType* ptrTy = dyn_cast<PointerType>(ptrOp->getType());
    IGC_ASSERT(ptrTy);
    Value* eltIdx = eltPtr->getOperand(1);

    builder->SetInsertPoint(eltPtr);
    // decide offset in bytes
    //     may need to create shift
    uint  eltBytes = int_cast<uint>(DL->getTypeStoreSize(ptrTy->getPointerElementType()));
    APInt eltSize = APInt(32, eltBytes);

    Value* offsetValue = eltIdx;
    if (eltSize != 1)
    {
        if (const ConstantInt * CI = dyn_cast<ConstantInt>(eltIdx))
        {
            uint32_t byteOffset = int_cast<uint32_t>(eltBytes * CI->getSExtValue());
            offsetValue = ConstantInt::get(eltIdx->getType(), byteOffset);
        }
        else if (eltSize.isPowerOf2())
        {
            APInt shift = APInt(32, eltSize.logBase2());
            offsetValue = builder->CreateShl(eltIdx, shift);
        }
        else
        {
            offsetValue = builder->CreateMul(eltIdx, builder->getInt(eltSize));
        }
    }
    return offsetValue;
}

Value* FixResourcePtr::CreateLoadIntrinsic(LoadInst* inst, Instruction* bufPtr, Value* offsetVal)
{
    Function* l;
    builder->SetInsertPoint(inst);
    llvm::Type* tys[2];
    tys[0] = inst->getType();
    tys[1] = bufPtr->getType();
    l = GenISAIntrinsic::getDeclaration(curFunc->getParent(),
        inst->getType()->isVectorTy() ? llvm::GenISAIntrinsic::GenISA_ldrawvector_indexed : llvm::GenISAIntrinsic::GenISA_ldraw_indexed,
        tys);

    unsigned alignment = std::max(inst->getType()->getScalarSizeInBits() / 8,
                                  (unsigned)inst->getAlignment());

    Value* attr[] =
    {
        bufPtr,
        offsetVal,
        builder->getInt32(alignment),
        builder->getInt1(inst->isVolatile())
    };
    Value* ld = builder->CreateCall(l, attr);
    if (!inst->getType()->isVectorTy())
    {
        if (!inst->getType()->isFloatTy())
        {
            Value* bitcast = dyn_cast<Instruction>(builder->CreateBitCast(ld, inst->getType()));
            ld = bitcast;
        }
    }
    return ld;
}

Value* FixResourcePtr::CreateStoreIntrinsic(StoreInst* inst, Instruction* bufPtr, Value* offsetVal)
{
    Function* l;
    builder->SetInsertPoint(inst);
    Value* storeVal = inst->getValueOperand();
    if (storeVal->getType()->isVectorTy())
    {
        llvm::Type* tys[2];
        tys[0] = bufPtr->getType();
        tys[1] = inst->getValueOperand()->getType();
        l = GenISAIntrinsic::getDeclaration(curFunc->getParent(),
            llvm::GenISAIntrinsic::GenISA_storerawvector_indexed,
            tys);
    }
    else
    {
        llvm::Type* dataType = storeVal->getType();

        IGC_ASSERT(dataType->getPrimitiveSizeInBits() == 16 || dataType->getPrimitiveSizeInBits() == 32);

        if (!dataType->isFloatingPointTy())
        {
            storeVal = builder->CreateBitCast(
                storeVal,
                dataType->getPrimitiveSizeInBits() == 32 ? builder->getFloatTy() : builder->getHalfTy());
        }

        llvm::Type* types[2] = {
            bufPtr->getType(),
            storeVal->getType() };

        l = GenISAIntrinsic::getDeclaration(curFunc->getParent(),
            llvm::GenISAIntrinsic::GenISA_storeraw_indexed,
            types);

    }
    unsigned alignment = std::max(storeVal->getType()->getScalarSizeInBits() / 8,
                                  (unsigned)inst->getAlignment());

    Value* attr[] =
    {
        bufPtr,
        offsetVal,
        storeVal,
        builder->getInt32(alignment),
        builder->getInt1(inst->isVolatile())
    };
    Value* st = builder->CreateCall(l, attr);
    return st;
}

///
/// @brief Resolves the buffer index argument of GetBufferPtr intrinsic.
///
/// This method looks for a ConstInt value of the buffer index. It handles
/// cases where the input buffer index is an extracted element of a vector,
/// e.g.
///
///  %17 = extractelement <2 x i32> zeroinitializer, i32 0
///  %19 = insertelement  <4 x i32> undef, i32 %17, i32 0
///  %20 = insertelement  <4 x i32> %19, i32 %15, i32 1
///  %21 = insertelement  <4 x i32> %20, i32 %15, i32 2
///  %Temp-47.i = insertelement <4 x i32> %21, i32 0, i32 3
///  %22 = extractelement <4 x i32> %Temp-47.i, i32 0
///  %23 = extractelement <4 x i32> %Temp-47.i, i32 1
///  %24 = extractelement <4 x i32> %Temp-47.i, i32 2
///  %25 = extractelement <4 x i32> %Temp-47.i, i32 3
///  %26 = call i32 addrspace(1179648)* @llvm.genx.GenISA.GetBufferPtr.p1179648i32(i32 %22, i32 1)
///
/// @param bufferIndex Buffer index value to resolve
/// @param vectorIndex Index inside of a vector where the buffer index is stored
/// @returns Value* ConstantInt buffer index if resolved or the input bufferIndex
///
Value* FixResourcePtr::ResolveBufferIndex(Value* bufferIndex, Value* vectorIndex)
{
    if (Constant * c = dyn_cast<Constant>(bufferIndex))
    {
        if (vectorIndex && isa<ConstantInt>(vectorIndex))
        {
            return c->getAggregateElement(cast<ConstantInt>(vectorIndex));
        }
        else if (isa<ConstantInt>(bufferIndex))
        {
            return bufferIndex;
        }
        else
        {
            IGC_ASSERT(0);
        }
    }
    else if (ExtractElementInst * ee = dyn_cast<ExtractElementInst>(bufferIndex))
    {
        return ResolveBufferIndex(
            ee->getOperand(0),  // vector argument to extract from
            ee->getOperand(1)); // element index in vector
    }
    else if (InsertElementInst * ie = dyn_cast<InsertElementInst>(bufferIndex))
    {
        if (vectorIndex &&
            isa<ConstantInt>(vectorIndex) &&
            isa<ConstantInt>(ie->getOperand(2)))
        {
            // compare indexes in vectors
            if ((cast<ConstantInt>(vectorIndex))->getZExtValue() == (cast<ConstantInt>(ie->getOperand(2)))->getZExtValue())
            {
                // indexes match,this is the insert element instruction with buffer index
                return ResolveBufferIndex(ie->getOperand(1));
            }
            else
            {
                // indexes don't match, keep looking for a matching insert element instruction
                return ResolveBufferIndex(ie->getOperand(0), vectorIndex);
            }
        }
    }

    return bufferIndex;
}