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

Copyright (C) 2017-2023 Intel Corporation

SPDX-License-Identifier: MIT

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

#pragma once

#include "Compiler/MetaDataUtilsWrapper.h"
#include "Compiler/CodeGenContextWrapper.hpp"
#include "Compiler/CISACodeGen/TranslationTable.hpp"

#include "common/LLVMWarningsPush.hpp"
#include <llvmWrapper/IR/InstrTypes.h>
#include <llvm/Pass.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/IntrinsicInst.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/ADT/Statistic.h>
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/DenseSet.h>
#include <llvm/ADT/SmallSet.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Value.h>
#include <llvm/IR/InstIterator.h>
#include <llvm/IR/Dominators.h>
#include <llvm/Analysis/LoopInfo.h>
#include <llvm/Analysis/PostDominators.h>
#include "common/LLVMWarningsPop.hpp"

#include <vector>
#include <common/Types.hpp>

namespace IGC
{
    class BranchInfo;
    class WIAnalysis;

    //This is a trick, since we cannot forward-declare enums embedded in class definitions.
    // The better solution is to completely hoist-out the WIDependency enum into a separate enum class
    // (c++ 11) and have it separate from WIAnalysis pass class. Nevertheless, that would require
    // updating many places in the current CodeGen code...
    // thus: WIAnalysis::WIDependancy ~ WIBaseClass::WIDependancy, the types are equal and do not require
    // conversion
    class WIBaseClass
    {
    public:
        /// @brief describes the type of dependency on the work item
        enum WIDependancy : uint8_t {
            UNIFORM_GLOBAL    = 0,  /// Same for all work-items within a shader.
            UNIFORM_WORKGROUP = 1,  /// Same for all work-items within a work group (compute).
            UNIFORM_THREAD    = 2,  /// Same for all work-items within a HW thread.
            CONSECUTIVE       = 3,  /// Elements are consecutive
            PTR_CONSECUTIVE   = 4,  /// Elements are pointers which are consecutive
            STRIDED           = 5,  /// Elements are in strides
            RANDOM            = 6,  /// Unknown or non consecutive order
            NumDeps           = 7,  /// Overall amount of dependencies
            INVALID           = 8
        };
    };

    // Provide FastValueMapAttributeInfo for WIDependancy.
    template<> struct FastValueMapAttributeInfo<WIBaseClass::WIDependancy> {
        static inline WIBaseClass::WIDependancy getEmptyAttribute() { return WIBaseClass::INVALID; }
    };

    class WIAnalysisRunner
    {
    public:
        void init(
            llvm::Function* F,
            llvm::LoopInfo* LI,
            llvm::DominatorTree* DT,
            llvm::PostDominatorTree* PDT,
            IGCMD::MetaDataUtils* MDUtils,
            CodeGenContext* CGCtx,
            ModuleMetaData* ModMD,
            TranslationTable* TransTable,
            bool ForCodegen = true);

        WIAnalysisRunner(
            llvm::Function* F,
            llvm::LoopInfo* LI,
            llvm::DominatorTree* DT,
            llvm::PostDominatorTree* PDT,
            IGCMD::MetaDataUtils* MDUtils,
            CodeGenContext* CGCtx,
            ModuleMetaData* ModMD,
            TranslationTable* TransTable,
            bool ForCodegen = true)
        {
            init(F, LI, DT, PDT, MDUtils, CGCtx, ModMD, TransTable, ForCodegen);
        }

        WIAnalysisRunner() {}

        bool run();

        /// @brief Returns the type of dependency the instruction has on
        /// the work-item
        /// @param val llvm::Value to test
        /// @return Dependency kind
        WIBaseClass::WIDependancy whichDepend(const llvm::Value* val) const;

        /// @brief Returns True if 'val' is uniform
        /// @param val llvm::Value to test
        bool isUniform(const llvm::Value* val) const;
        bool isWorkGroupOrGlobalUniform(const llvm::Value* val) const;
        bool isGlobalUniform(const llvm::Value* val) const;

        /// incremental update of the dep-map on individual value
        /// without propagation. Exposed for later pass.
        void incUpdateDepend(const llvm::Value* val, WIBaseClass::WIDependancy dep)
        {
            m_depMap.SetAttribute(val, dep);
        }

        /// check if a value is defined inside divergent control-flow
        bool insideDivergentCF(const llvm::Value* val) const
        {
            return(llvm::isa<llvm::Instruction>(val) &&
                m_ctrlBranches.find(llvm::cast<llvm::Instruction>(val)->getParent()) != m_ctrlBranches.end());
        }

        /// check if a value is defined inside workgroup divergent control-flow.
        /// This will return false if the value is in the influence region
        /// of only global and workgroup uniform branches.
        bool insideWorkgroupDivergentCF(const llvm::Value* val) const;

        void releaseMemory()
        {
            m_ctrlBranches.clear();
            m_changed1.clear();
            m_changed2.clear();
            m_allocaDepMap.clear();
            m_storeDepMap.clear();
            m_depMap.clear();
            m_forcedUniforms.clear();
        }

        /// print - print m_deps in human readable form
        void print(llvm::raw_ostream& OS, const llvm::Module* = 0) const;

        /// dump - Dump the m_deps to a file.
        void dump() const;

        // helper for dumping WI info into files with lock
        void lock_print();

    private:
        WIBaseClass::WIDependancy getCFDependency(const llvm::BasicBlock* BB) const;

        struct AllocaDep
        {
            std::vector<const llvm::StoreInst*> stores;
            std::vector<const llvm::IntrinsicInst*> lifetimes;
            bool assume_uniform = false;
        };

        /// @brief Update dependency relations between all values
        void updateDeps();

        /// @brief mark the arguments dependency based on the metadata set
        void updateArgsDependency(llvm::Function* pF);

        /*! \name Dependency Calculation Functions
         *  \{ */
         /// @brief Calculate the dependency type for the instruction
         /// @param inst Instruction to inspect
         /// @return Type of dependency.
        void calculate_dep(const llvm::Value* val);
        WIBaseClass::WIDependancy calculate_dep(const llvm::BinaryOperator* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::CallInst* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::GetElementPtrInst* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::PHINode* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::SelectInst* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::AllocaInst* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::CastInst* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::VAArgInst* inst);
        WIBaseClass::WIDependancy calculate_dep(const llvm::LoadInst* inst);

        WIBaseClass::WIDependancy calculate_dep_terminator(const IGCLLVM::TerminatorInst* inst);
        /*! \} */

        /// @brief do the trivial checking WI-dep
        /// @param I instruction to check
        /// @return Dependency type. Returns Uniform if all operands are
        ///         Uniform, Random otherwise
        WIBaseClass::WIDependancy calculate_dep_simple(const llvm::Instruction* I);

        /// @brief update the WI-dep from a divergent branch,
        ///        affected instructions are added to m_pChangedNew
        /// @param the divergent branch
        void update_cf_dep(const IGCLLVM::TerminatorInst* TI);

        /// @brief update the WI-dep for a sequence of insert-elements forming a vector
        ///        affected instructions are added to m_pChangedNew
        /// @param the insert-element instruction
        void updateInsertElements(const llvm::InsertElementInst* inst);

        /// @brief update the WI-dep for a insertValue chain to add affected
        ///        instructions to m_pChangedNew. This is to make sure if
        ///        one in the chain is RANDOM, all are RANDOM.
        /// @param the insert-element instruction
        void updateInsertValues(const llvm::InsertValueInst* Inst);

        /// @check phi divergence at a join-blk due to a divergent branch
        void updatePHIDepAtJoin(llvm::BasicBlock* blk, BranchInfo* brInfo);

        void updateDepMap(const llvm::Instruction* inst, WIBaseClass::WIDependancy dep);

        /// @brief Provide known dependency type for requested value
        /// @param val llvm::Value to examine
        /// @return Dependency type. Returns Uniform for unknown type
        WIBaseClass::WIDependancy getDependency(const llvm::Value* val);

        /// @brief return true if there is calculated dependency type for requested value
        /// @param val llvm::Value to examine
        /// @return true if value has dependency type, false otherwise.
        bool hasDependency(const llvm::Value* val) const;

        /// @brief return true if all uses of this value are marked RANDOM
        bool allUsesRandom(const llvm::Value* val);

        /// @brief return true if any of the use require the value to be uniform
        bool needToBeUniform(const llvm::Value* val);

        /// @brief return true is the instruction is simple and making it random is cheap
        bool isInstructionSimple(const llvm::Instruction* inst);

        /// @brief return true if all the source operands are defined outside the region
        bool isRegionInvariant(const llvm::Instruction* inst, BranchInfo* brInfo);

        /// @brief return true if instruction is used as lane ID in subgroup broadcast
        bool isUsedByWaveBroadcastAsLocalID(const llvm::Instruction* inst);

        /// @brief update dependency structure for Alloca
        bool TrackAllocaDep(const llvm::Value* I, AllocaDep& dep);

        void checkLocalIdUniform(
            llvm::Function* F,
            bool& IsLxUniform,
            bool& IsLyUniform,
            bool& IsLzUniform);

        void CS_checkLocalIDs(llvm::Function *F);
    private:

        /// The WIAnalysis follows pointer arithmetic
        ///  and Index arithmetic when calculating dependency
        ///  properties. If a part of the index is lost due to
        ///  a transformation, it is acceptable.
        ///  This constant decides how many bits need to be
        ///  preserved before we give up on the analysis.
        static const unsigned int MinIndexBitwidthToPreserve;

        /// Stores an updated list of all dependencies
        /// for each block, store the list of diverging branches that affect it
        llvm::DenseMap<const llvm::BasicBlock*, llvm::SmallPtrSet<const llvm::Instruction*, 4>> m_ctrlBranches;

        /// Iteratively one set holds the changed from the previous iteration and
        /// the other holds the new changed values from the current iteration.
        std::vector<const llvm::Value*> m_changed1{};
        std::vector<const llvm::Value*> m_changed2{};
        /// ptr to m_changed1, m_changed2
        std::vector<const llvm::Value*>* m_pChangedOld = nullptr;
        std::vector<const llvm::Value*>* m_pChangedNew = nullptr;

        /// <summary>
        ///  hold the vector-defs that are promoted from an uniform alloca
        ///  therefore, need to be forced into uniform no matter what.
        /// </summary>
        std::vector<const llvm::Value*> m_forcedUniforms{};

        llvm::Function* m_func = nullptr;
        llvm::LoopInfo *LI = nullptr;
        llvm::DominatorTree* DT = nullptr;
        llvm::PostDominatorTree* PDT = nullptr;
        IGC::IGCMD::MetaDataUtils* m_pMdUtils = nullptr;
        IGC::CodeGenContext* m_CGCtx = nullptr;
        IGC::ModuleMetaData* m_ModMD = nullptr;
        IGC::TranslationTable* m_TT = nullptr;

        // Allow access to all the store into an alloca if we were able to track it
        llvm::DenseMap<const llvm::AllocaInst*, AllocaDep> m_allocaDepMap;
        // reverse map to allow to know what alloca to update when store changes
        llvm::DenseMap<const llvm::StoreInst*, const llvm::AllocaInst*> m_storeDepMap;

        IGC::FastValueMap<WIBaseClass::WIDependancy, FastValueMapAttributeInfo<WIBaseClass::WIDependancy>> m_depMap;

        // For dumpping WIA info per each invocation
        static llvm::DenseMap<const llvm::Function*, int> m_funcInvocationId;

        // Is this analysis providing information that could be used for late
        // stage codegen, or is it just used to determine uniformity early on?
        bool ForCodegen = true;

        // Caching CS local ID's uniformness
        bool m_localIDxUniform = false;
        bool m_localIDyUniform = false;
        bool m_localIDzUniform = false;
    };

    /// @brief Work Item Analysis class used to provide information on
    ///  individual instructions. The analysis class detects values which
    ///  depend in work-item and describe their dependency.
    ///  The algorithm used is recursive and new instructions are updated
    ///  according to their operands (which are already calculated).
    ///  original code for OCL vectorizer
    ///
    class WIAnalysis : public llvm::FunctionPass, public WIBaseClass
    {
    public:
        static char ID; // Pass identification, replacement for typeid

        WIAnalysis();

        ~WIAnalysis() {}

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

        void getAnalysisUsage(llvm::AnalysisUsage& AU) const override
        {
            // Analysis pass preserve all
            AU.setPreservesAll();
            AU.addRequired<llvm::DominatorTreeWrapperPass>();
            AU.addRequired<llvm::PostDominatorTreeWrapperPass>();
            AU.addRequired<llvm::LoopInfoWrapperPass>();
            AU.addRequired<MetaDataUtilsWrapper>();
            AU.addRequired<CodeGenContextWrapper>();
            AU.addRequired<TranslationTable>();
        }

        /// @brief LLVM llvm::Function pass entry
        /// @param F llvm::Function to transform
        /// @return True if changed
        bool runOnFunction(llvm::Function& F) override;

        /// print - print m_deps in human readable form
        void print(llvm::raw_ostream& OS, const llvm::Module* = 0) const override;

        /// dump - Dump the m_deps to dbgs().
        void dump() const;

    public:
        /// @brief Returns the type of dependency the instruction has on
        /// the work-item
        /// @param val llvm::Value to test
        /// @return Dependency kind
        WIDependancy whichDepend(const llvm::Value* val);

        /// @brief Returns True if 'val' is uniform
        /// @param val llvm::Value to test
        bool isUniform(const llvm::Value* val) const; // Return true for any uniform
        bool isWorkGroupOrGlobalUniform(const llvm::Value* val);
        bool isGlobalUniform(const llvm::Value* val);

        /// incremental update of the dep-map on individual value
        /// without propagation. Exposed for later pass.
        void incUpdateDepend(const llvm::Value* val, WIDependancy dep);

        /// check if a value is defined inside divergent control-flow
        bool insideDivergentCF(const llvm::Value* val) const;

        /// check if a value is defined inside workgroup divergent control-flow
        /// This will return false if the value is in the influence region
        /// of only global and workgroup uniform branches.
        bool insideWorkgroupDivergentCF(const llvm::Value* val) const;

        void releaseMemory() override
        {
            Runner.releaseMemory();
        }


        /// Return true if Dep is any of uniform dependancy.
        static bool isDepUniform(WIDependancy Dep) {
            return Dep == WIDependancy::UNIFORM_GLOBAL ||
                Dep == WIDependancy::UNIFORM_WORKGROUP ||
                Dep == WIDependancy::UNIFORM_THREAD;
        }
        WIAnalysisRunner Runner;
    };

} // namespace IGC

