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

Copyright (C) 2020-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "AnnotateUniformAllocas.h"
#include "GenISAIntrinsics/GenIntrinsicInst.h"
#include "Compiler/IGCPassSupport.h"
#include "IGCIRBuilder.h"
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Function.h>
#include <llvm/Transforms/Utils/Local.h>
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;
using namespace IGC;

#define PASS_FLAG "annotate_uniform_allocas"
#define PASS_DESCRIPTION "Annotate uniform allocas"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(AnnotateUniformAllocas, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_END(AnnotateUniformAllocas, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

namespace IGC
{
    char AnnotateUniformAllocas::ID = 0;

    AnnotateUniformAllocas::AnnotateUniformAllocas() : FunctionPass(ID)
    {
        initializeAnnotateUniformAllocasPass(*PassRegistry::getPassRegistry());
    }

    llvm::FunctionPass* createAnnotateUniformAllocasPass()
    {
        return new AnnotateUniformAllocas();
    }

    bool AnnotateUniformAllocas::runOnFunction(Function& F)
    {
        WI = &getAnalysis<WIAnalysis>();
        IGC_ASSERT(WI != nullptr);
        visit(F);

        while (!AssumeToErase.empty()) {
            auto Inst = AssumeToErase.pop_back_val();
            bool IsAssume = false;
            if (llvm::GenIntrinsicInst* pIntr = llvm::dyn_cast<llvm::GenIntrinsicInst>(Inst)) {
                IsAssume = (pIntr->getIntrinsicID() == GenISAIntrinsic::GenISA_assume_uniform);
            }
            if (IsAssume || isInstructionTriviallyDead(Inst)) {
                for (Use& OpU : Inst->operands()) {
                    Value* OpV = OpU.get();
                    if (Instruction* OpI = dyn_cast<Instruction>(OpV))
                        AssumeToErase.push_back(OpI);
                }
                Inst->eraseFromParent();
            }
        }
        return m_changed;
    }

    void AnnotateUniformAllocas::visitAllocaInst(AllocaInst& I)
    {
        bool isUniformAlloca = WI->isUniform(&I);
        if (isUniformAlloca)
        {
            // add the meta-date to the alloca for promotion
            IRBuilder<> builder(&I);
            MDNode* node = MDNode::get(I.getContext(), ConstantAsMetadata::get(builder.getInt1(true)));
            I.setMetadata("uniform", node);
            m_changed = true;
        }
    }

    void AnnotateUniformAllocas::visitCallInst(CallInst& I)
    {
        if (llvm::GenIntrinsicInst* pIntr = llvm::dyn_cast<llvm::GenIntrinsicInst>(&I))
        {
            if (pIntr->getIntrinsicID() == GenISAIntrinsic::GenISA_assume_uniform)
            {
                AssumeToErase.push_back(pIntr);
                // add the meta-date to the alloca for promotion
                auto OpV = pIntr->getOperand(0);
                while (OpV)
                {
                   if (auto CI = dyn_cast<CastInst>(OpV))
                   {
                       OpV = CI->getOperand(0);
                       continue;
                   }
                   if (auto ALI = dyn_cast<AllocaInst>(OpV))
                   {
                       IRBuilder<> builder(ALI);
                       MDNode* node = MDNode::get(ALI->getContext(),
                               ConstantAsMetadata::get(builder.getInt1(true)));
                       ALI->setMetadata("UseAssumeUniform", node);
                   }
                   break;
                }
            }
        }
    }
}
