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

Copyright (c) 2017 Intel Corporation

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.


======================= end_copyright_notice ==================================*/
#include "Compiler/CISACodeGen/LinkTessControlShaderPass.h"
#include "Compiler/IGCPassSupport.h"
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Constants.h>
#include <llvm/IR/Intrinsics.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/InstVisitor.h>
#include <llvm/IR/IRBuilder.h>
#include "common/LLVMWarningsPop.hpp"

#include "Compiler/MetaDataUtilsWrapper.h"
#include "Compiler/CodeGenPublic.h"

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

// Summary: The role of this pass is to link the TCS generated by OGL FE 
// by adding a loop to loop through the number of output control points and 
// replace all occurances of HSControlPointID with the loopCounter
namespace IGC
{
    class LinkTessControlShader : public llvm::ModulePass
    {
    public:
        // Pass identification, replacement for typeid
        static char ID;

        /// @brief  Constructor
        LinkTessControlShader();
        ~LinkTessControlShader() {};

        /// @brief  Provides name of pass
        virtual llvm::StringRef getPassName() const override {
            return "LinkTessControlShader";
        }

        virtual void getAnalysisUsage(llvm::AnalysisUsage &AU) const override
        {
            AU.addRequired<MetaDataUtilsWrapper>();
            AU.addRequired<CodeGenContextWrapper>();
        }

        /// @brief  Main entry point.
        /// @param  F The current function.
        virtual bool runOnModule(llvm::Module &M) override;
    private:
        HullShaderDispatchModes DetermineDispatchMode(llvm::Module* mod, llvm::IRBuilder<> *irBuilder);
        static const uint32_t SIMDSize;
        llvm::Function* CreateNewTCSFunction(llvm::Function* pCurrentFunc);
        bool CheckIfBarrierInstructionExists(llvm::Function* pFunc);
        bool m_useMultipleHardwareThread;
    };

    char LinkTessControlShader::ID = 0;
    // Register pass to igc-opt
    #define PASS_FLAG "igc-linkTessControlShader"
    #define PASS_DESCRIPTION "Perform looping of tessellation function based on control point count"
    #define PASS_CFG_ONLY false
    #define PASS_ANALYSIS false
    IGC_INITIALIZE_PASS_BEGIN(LinkTessControlShader, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
    IGC_INITIALIZE_PASS_END(LinkTessControlShader, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

    // SIMD size for tessellation workloads is SIMD8
    const uint32_t LinkTessControlShader::SIMDSize = 8;

    LinkTessControlShader::LinkTessControlShader() : llvm::ModulePass(ID), m_useMultipleHardwareThread(false)
    {
        initializeLinkTessControlShaderPass(*llvm::PassRegistry::getPassRegistry());
    }

    HullShaderDispatchModes LinkTessControlShader::DetermineDispatchMode(llvm::Module* mod, llvm::IRBuilder<> *irBuilder)
    {
        llvm::NamedMDNode *metaData = mod->getOrInsertNamedMetadata("HullShaderDispatchMode");

        // now find out the output control point count
        llvm::GlobalVariable* pOCPCount = mod->getGlobalVariable("HSOutputControlPointCount");
        uint32_t outputControlPointCount = int_cast<uint32_t>(llvm::cast<llvm::ConstantInt>(pOCPCount->getInitializer())->getZExtValue());

        IGC::CodeGenContext* pCodeGenContext = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();

        /* Instance Count
        **         This field determines the number of threads(minus one) spawned per input patch.

        **         If the HS kernel uses a barrier function, software must restrict the Instance Count 
        **         to the number of threads that can be simultaneously active within a subslice.Factors 
        **         which must be considered includes scratch memory availability.
        **         Value	     	Description
        **         [0, 15]	 	    representing[1, 16] instances */

        llvm::GlobalVariable* pGlobal = mod->getGlobalVariable("TessInputControlPointCount");
        unsigned int inputControlPointCount = int_cast<unsigned int>(llvm::cast<llvm::ConstantInt>(pGlobal->getInitializer())->getZExtValue());

        //use HS single patch if WA exists and input control points >= 29 as not enough registers for push constants
        bool useSinglePatch = !(pCodeGenContext->platform.WaDispatchGRFHWIssueInGSAndHSUnit() && inputControlPointCount >= 29);

        if (pCodeGenContext->platform.useOnlyEightPatchDispatchHS() ||
            (pCodeGenContext->platform.supportHSEightPatchDispatch() &&
             !(m_useMultipleHardwareThread && outputControlPointCount >= 16) &&
             useSinglePatch &&
             IGC_IS_FLAG_DISABLED(EnableHSSinglePatchDispatch)))
        {
            Constant* cval = llvm::ConstantInt::get(
                irBuilder->getInt32Ty(),
                HullShaderDispatchModes::EIGHT_PATCH_DISPATCH_MODE);
            llvm::MDNode *mdNode = llvm::MDNode::get(
                irBuilder->getContext(),
                llvm::ConstantAsMetadata::get(cval));
            metaData->addOperand(mdNode);
            return HullShaderDispatchModes::EIGHT_PATCH_DISPATCH_MODE;
        }
        else
        {
            Constant* cval = llvm::ConstantInt::get(
                irBuilder->getInt32Ty(),
                HullShaderDispatchModes::SINGLE_PATCH_DISPATCH_MODE);
            llvm::MDNode *mdNode = llvm::MDNode::get(
                irBuilder->getContext(),
                llvm::ConstantAsMetadata::get(cval));
            metaData->addOperand(mdNode);
            return HullShaderDispatchModes::SINGLE_PATCH_DISPATCH_MODE;
        }
    }
    /*
        Function pass to create a new entry function
    */
    llvm::Function* LinkTessControlShader::CreateNewTCSFunction(llvm::Function* pCurrentFunc)
    {
        llvm::IRBuilder<> irBuilder(pCurrentFunc->getContext());
        CodeGenContext *ctx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();

        std::vector<llvm::Type*> callArgTypes;
        for (auto& argIter : range(pCurrentFunc->arg_begin(), pCurrentFunc->arg_end()))
        {
            callArgTypes.push_back(argIter.getType());
        }
        callArgTypes.push_back(irBuilder.getInt32Ty());

        std::string funcName = "tessControlShaderEntry";

        llvm::Function* pNewFunction = llvm::Function::Create(
            llvm::FunctionType::get(
                irBuilder.getVoidTy(), callArgTypes, false),
            llvm::GlobalValue::PrivateLinkage,
            funcName,
            ctx->getModule());

        pNewFunction->addFnAttr(llvm::Attribute::AlwaysInline);

        // Move over the contents of the original function
        pNewFunction->getBasicBlockList().splice(pNewFunction->begin(), pCurrentFunc->getBasicBlockList());

        llvm::Function* pToBeDeletedFunc = pCurrentFunc;

        for (auto oldArg = pToBeDeletedFunc->arg_begin(),
            oldArgEnd = pToBeDeletedFunc->arg_end(),
            newArg = pNewFunction->arg_begin();
        oldArg != oldArgEnd;
        ++oldArg, ++newArg)
        {
            oldArg->replaceAllUsesWith(&(*newArg));
            newArg->takeName(&(*oldArg));
        }

        // delete the old function signature
        pToBeDeletedFunc->eraseFromParent();

        // now replace all occurences of HSControlPointID with the current 
        // loop iteration CPID - pCurrentCPID
       SmallVector<Instruction*, 10> instructionToRemove;

       llvm::Value* pHSControlPointID = llvm::GenISAIntrinsic::getDeclaration(pNewFunction->getParent(),
            GenISAIntrinsic::GenISA_DCL_HSControlPointID);

       unsigned int argIndexInFunc = pNewFunction->getArgumentList().size()-1;
       Function::arg_iterator arg = pNewFunction->arg_begin();
       for (unsigned int i = 0; i < argIndexInFunc; ++i, ++arg);

       for (Value::user_iterator i = pHSControlPointID->user_begin(), e = pHSControlPointID->user_end(); i != e; ++i)
        {
            Instruction *useInst = cast<Instruction>(*i);
            if (useInst->getParent()->getParent() == pNewFunction)
            {
                instructionToRemove.push_back(useInst);
                useInst->replaceAllUsesWith(&(*arg));
            }
        }

        for (auto &inst : instructionToRemove)
        {
            inst->eraseFromParent();
        }
        return pNewFunction;
    }

    bool LinkTessControlShader::CheckIfBarrierInstructionExists(llvm::Function* pFunc)
    {
        for (llvm::Function::iterator bb = pFunc->begin(), be = pFunc->end();
            bb != be;
            ++bb)
        {
            for (llvm::BasicBlock::iterator ii = bb->begin(), ie = bb->end();
                ii != ie;
                ++ii)
            {
                if (GenIntrinsicInst* inst = dyn_cast<GenIntrinsicInst>(ii))
                {
                    if (inst->getIntrinsicID() == GenISAIntrinsic::GenISA_threadgroupbarrier)
                    {
                        return true;
                    }
                }
            }
        }

        return false;
    }

    bool LinkTessControlShader::runOnModule(llvm::Module &M)
    {
        CodeGenContext *ctx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
        MetaDataUtils *pMdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
        if (pMdUtils->size_FunctionsInfo() != 1)
        {
            return false;
        }

        Function* pFunc = pMdUtils->begin_FunctionsInfo()->first;
        llvm::IRBuilder<> builder(pFunc->getContext());
        std::string oldEntryFuncName = pFunc->getName().str();

        llvm::Function* pNewTCSFunction = CreateNewTCSFunction(pFunc);

        m_useMultipleHardwareThread = CheckIfBarrierInstructionExists(pNewTCSFunction);

        // Determine the dispatch mode
        HullShaderDispatchModes dispatchMode = DetermineDispatchMode(pNewTCSFunction->getParent(), &builder);

        // This function is the new entry function
        llvm::Function* pNewLoopFunc = llvm::Function::Create(llvm::FunctionType::get(builder.getVoidTy(), false),
            llvm::GlobalValue::ExternalLinkage,
            oldEntryFuncName,
            ctx->getModule());

        llvm::BasicBlock* pEntryBlock = llvm::BasicBlock::Create(
            pNewLoopFunc->getContext(),
            oldEntryFuncName,
            pNewLoopFunc);

        builder.SetInsertPoint(pEntryBlock);

        // first create a call to simdLaneId() intrinsic
        llvm::Value* pCPId = nullptr;
        llvm::Function *pFuncPatchInstanceIdOrSIMDLaneId = nullptr;
        switch (dispatchMode)
        {
        case SINGLE_PATCH_DISPATCH_MODE:
            pFuncPatchInstanceIdOrSIMDLaneId = llvm::GenISAIntrinsic::getDeclaration(
                pNewLoopFunc->getParent(), llvm::GenISAIntrinsic::GenISA_simdLaneId);
            pCPId = builder.CreateCall(pFuncPatchInstanceIdOrSIMDLaneId);

            if (m_useMultipleHardwareThread)
            {
 				// CPID = patchInstanceID * 8 + SimdLaneId;
                pFuncPatchInstanceIdOrSIMDLaneId = llvm::GenISAIntrinsic::getDeclaration(
                    pNewLoopFunc->getParent(), llvm::GenISAIntrinsic::GenISA_patchInstanceId);
                pCPId = builder.CreateAdd(
                    builder.CreateZExt(
                        pCPId,
                        builder.getInt32Ty()),
                    builder.CreateMul(
                        builder.CreateCall(pFuncPatchInstanceIdOrSIMDLaneId),
                        llvm::ConstantInt::get(builder.getInt32Ty(), SIMDSize)));
            }
            break;

        case EIGHT_PATCH_DISPATCH_MODE:
            if(m_useMultipleHardwareThread)
            {
                pFuncPatchInstanceIdOrSIMDLaneId = llvm::GenISAIntrinsic::getDeclaration(
                    pNewLoopFunc->getParent(), llvm::GenISAIntrinsic::GenISA_patchInstanceId);
                pCPId = builder.CreateCall(pFuncPatchInstanceIdOrSIMDLaneId);
            }
            else
            { 
                pCPId = builder.getInt32(0);
            }
            break;

        default:
            assert(0 && "should not reach here");
            break;
        }

        // now find out the output control point count
        llvm::GlobalVariable* pOCPCount = pNewLoopFunc->getParent()->getGlobalVariable("HSOutputControlPointCount");
        uint32_t outputControlPointCount = int_cast<uint32_t>(llvm::cast<llvm::ConstantInt>(pOCPCount->getInitializer())->getZExtValue());

        // We don't need to deal with any loops when we are using multiple hardware threads
        if (!m_useMultipleHardwareThread)
        {
            // initialize instanceCount to output control point count
            llvm::Value* pInstanceCount = builder.getInt32(outputControlPointCount);

            // initialize loopCounter
            llvm::Value* pLoopCounter = builder.CreateAlloca(builder.getInt32Ty(), 0, "loopCounter");
            llvm::Value* pConstInt = builder.getInt32(0);
            builder.CreateStore(pConstInt, pLoopCounter, false);

            // create loop-entry basic block and setInsertPoint to loop-entry
            //llvm::Function* pParentFunc = builder.GetInsertBlock()->getParent();
            llvm::BasicBlock* pLoopEntryBB = llvm::BasicBlock::Create(pNewLoopFunc->getContext(),
                VALUE_NAME("tcs_loop_entry"),
                pNewLoopFunc);

            llvm::BasicBlock* pLoopConditionTrue = llvm::BasicBlock::Create(pNewLoopFunc->getContext(),
                VALUE_NAME("tcs_loop_condition_true"),
                pNewLoopFunc);

            // Create loop-continue basic block
            llvm::BasicBlock* pLoopContinueBB = llvm::BasicBlock::Create(pNewLoopFunc->getContext(),
                VALUE_NAME("tcs_loop_continue"),
                pNewLoopFunc);

            // create loop exit basic block
            llvm::BasicBlock* pAfterLoopBB = llvm::BasicBlock::Create(pNewLoopFunc->getContext(),
                VALUE_NAME("tcs_after_loop"),
                pNewLoopFunc);

            // setInsertPoint to loopEntryBB
            builder.CreateBr(pLoopEntryBB);
            builder.SetInsertPoint(pLoopEntryBB);

            // Load the loop counter
            llvm::LoadInst* pLoadLoopCounter = builder.CreateLoad(pLoopCounter);
            llvm::Value* pMulLoopCounterRes = nullptr;
            llvm::Value* pCurrentCPID = nullptr;
            llvm::Value* pConditionalRes1 = nullptr;
            uint32_t loopIterationCount = 0;

            switch (dispatchMode)
            {
            case SINGLE_PATCH_DISPATCH_MODE:
                // currentCPID = pCPId + loopCounter x simdsize ( in this case its always simd 8 )
                pMulLoopCounterRes = builder.CreateMul(
                    pLoadLoopCounter,
                    llvm::ConstantInt::get(builder.getInt32Ty(), SIMDSize));
                pCurrentCPID = builder.CreateAdd(
                    builder.CreateZExt(
                    pCPId,
                    builder.getInt32Ty()),
                    pMulLoopCounterRes);

                // cmp currentCPID to instanceCount so we enable only the required lanes
                pConditionalRes1 = builder.CreateICmpULT(
                    pCurrentCPID,
                    pInstanceCount,
                    VALUE_NAME("tcs_if_ult_cond1"));

                // if true go to startBB else jump to pAfterLoopBB
                builder.CreateCondBr(pConditionalRes1,
                    pLoopConditionTrue,
                    pAfterLoopBB);

                loopIterationCount = ((outputControlPointCount - 1) / 8) + 1;
                break;

            case EIGHT_PATCH_DISPATCH_MODE:
                pCurrentCPID = pLoadLoopCounter;
                loopIterationCount = outputControlPointCount;

                // jump to startBB
                builder.CreateBr(pLoopConditionTrue);
                break;

            default:
                assert(false && "should not reach here");
                break;
            }

            // branch to pLoopContinueBB from endBB
            builder.SetInsertPoint(pLoopConditionTrue);

            // Create a call to the TCS function when condition is true to loop the function as many times as the number of control points
            builder.CreateCall(pNewTCSFunction, pCurrentCPID);
            builder.CreateBr(pLoopContinueBB);

            // setInsertPoint to pLoopContinueBB
            builder.SetInsertPoint(pLoopContinueBB);
            // increment loop counter loopCounter = loopCounter + 1
            llvm::Value* pIncrementedLoopCounter = builder.CreateAdd(
                pLoadLoopCounter,
                llvm::ConstantInt::get(builder.getInt32Ty(), 1));
            builder.CreateStore(pIncrementedLoopCounter, pLoopCounter, false);

            // now evaluate loop, if( ( incrementedLoopCounter ) < ( ( maxControlPointCount - 1 )/8) + 1 )
            // then continue loop else go to after loop
            llvm::Value* pNumberOfLoopIterationsRequired = llvm::ConstantInt::get(builder.getInt32Ty(), loopIterationCount);

            llvm::Value* pConditionalRes2 = builder.CreateICmpULT(
                pIncrementedLoopCounter,
                pNumberOfLoopIterationsRequired,
                VALUE_NAME("tcs_if_ult_cond2"));

            // create branch to LoopEntryBB or AfterLoopBB based on result of conditional branch
            builder.CreateCondBr(pConditionalRes2,
                pLoopEntryBB,
                pAfterLoopBB);

            // set insert point to afterloop basic block
            builder.SetInsertPoint(pAfterLoopBB);
        }
        else if (dispatchMode == SINGLE_PATCH_DISPATCH_MODE)
        {
            // In single patch dispatch mode the execution mask is 0xFF. Make 
            // that only valid CPIDs execute.

            // Create the main basic block for the shader
            llvm::BasicBlock* pTcsBody = llvm::BasicBlock::Create(pNewLoopFunc->getContext(),
                VALUE_NAME("tcs_body"),
                pNewLoopFunc);
            // and the end block.
            llvm::BasicBlock* pAfterTcsBody = llvm::BasicBlock::Create(pNewLoopFunc->getContext(),
                VALUE_NAME("tcs_end"),
                pNewLoopFunc);

            // Compare current CPID to the number of CPIDs to enable only the required lanes.
            llvm::Value* pIsLaneEnabled = builder.CreateICmpULT(
                pCPId,
                builder.getInt32(outputControlPointCount),
                VALUE_NAME("tcs_if_ult_cond1"));

            builder.CreateCondBr(pIsLaneEnabled,
                pTcsBody,
                pAfterTcsBody);

            builder.SetInsertPoint(pTcsBody);

            // Call TCS function.
            builder.CreateCall(pNewTCSFunction, pCPId);

            builder.CreateBr(pAfterTcsBody);
            builder.SetInsertPoint(pAfterTcsBody);
        }
        else
        {
            // when using multiple hardware threads just call the Control Point function once with the appropriate CPID
            builder.CreateCall(pNewTCSFunction, pCPId);
        }

        // add terminator to the afterloop basic block
        builder.CreateRetVoid();

        pMdUtils->clearFunctionsInfo();
        IGCMetaDataHelper::addFunction(*pMdUtils, pNewLoopFunc);
        return true;
    }

    llvm::Pass* createLinkTessControlShader()
    {
        return new LinkTessControlShader();
    }
}
