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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#pragma once

#include "Compiler/CISACodeGen/CISACodeGen.h"
#include "Compiler/CISACodeGen/WIAnalysis.hpp"
#include "Compiler/MetaDataApi/MetaDataApi.h"
#include "GenISAIntrinsics/GenIntrinsics.h"

#include <vector>
#include <set>
#include <map>
#include <string>

namespace llvm
{
    class LLVMContext;
    class Function;
    class Type;
}

namespace IGC
{
    /// @brief  ImplicitArg is used for representing the implict information that is passed from the
    ///         OpenCL runtime to IGC, characterize it and allow IGC to use it.
    /// @Author Marina Yatsina
    class ImplicitArg {
    public:

        /// @brief  Type of implicit information passed from the OpenCL runtime
        enum ArgType {
            R0,
            START_ID = R0,
            PAYLOAD_HEADER,

            // WI information
            WORK_DIM,
            NUM_GROUPS,
            GLOBAL_SIZE,
            LOCAL_SIZE,
            ENQUEUED_LOCAL_WORK_SIZE,
            LOCAL_ID_X,
            LOCAL_ID_Y,
            LOCAL_ID_Z,

            // Pointer bases and buffers
            CONSTANT_BASE,
            GLOBAL_BASE,
            PRIVATE_BASE,
            PRINTF_BUFFER,

            // Buffer offset (for stateless to stateful optim)
            BUFFER_OFFSET,

            // Aggregates
            STRUCT_START,
            CONSTANT_REG_FP32 = STRUCT_START,
            CONSTANT_REG_QWORD,
            CONSTANT_REG_DWORD,
            CONSTANT_REG_WORD,
            CONSTANT_REG_BYTE,
            STRUCT_END = CONSTANT_REG_BYTE,

            // Images
            IMAGES_START,
            IMAGE_HEIGHT = IMAGES_START,
            IMAGE_WIDTH,
            IMAGE_DEPTH,
            IMAGE_NUM_MIP_LEVELS,
            IMAGE_CHANNEL_DATA_TYPE,
            IMAGE_CHANNEL_ORDER,
            IMAGE_SRGB_CHANNEL_ORDER,
            IMAGE_ARRAY_SIZE,
            IMAGE_NUM_SAMPLES,
            SAMPLER_ADDRESS,
            SAMPLER_NORMALIZED,
            SAMPLER_SNAP_WA,
            INLINE_SAMPLER,
            FLAT_IMAGE_BASEOFFSET,
            FLAT_IMAGE_HEIGHT,
            FLAT_IMAGE_WIDTH,
            FLAT_IMAGE_PITCH,
            IMAGES_END = FLAT_IMAGE_PITCH,

            // VME
            VME_MB_BLOCK_TYPE,
            VME_SUBPIXEL_MODE,
            VME_SAD_ADJUST_MODE,
            VME_SEARCH_PATH_TYPE,

            // Device Enqueue
            DEVICE_ENQUEUE_DEFAULT_DEVICE_QUEUE,
            DEVICE_ENQUEUE_EVENT_POOL,
            DEVICE_ENQUEUE_MAX_WORKGROUP_SIZE,
            DEVICE_ENQUEUE_PARENT_EVENT,
            DEVICE_ENQUEUE_PREFERED_WORKGROUP_MULTIPLE,
            GET_OBJECT_ID,
            GET_BLOCK_SIMD_SIZE,

            // GAS
            LOCAL_MEMORY_STATELESS_WINDOW_START_ADDRESS,
            LOCAL_MEMORY_STATELESS_WINDOW_SIZE,
            PRIVATE_MEMORY_STATELESS_SIZE,

            STAGE_IN_GRID_ORIGIN,
            STAGE_IN_GRID_SIZE,

            SYNC_BUFFER,

            // Raytracing args
            RT_GLOBAL_BUFFER_POINTER,
            RT_LOCAL_BUFFER_POINTER,
            RT_INLINED_DATA,
            RT_STACK_ID,

            // Bindless buffer (for stateless to bindless optim)
            BINDLESS_OFFSET,

            // Pointer to implicit arguments prepared by runtime
            IMPLICIT_ARG_BUFFER_PTR,

            ASSERT_BUFFER_POINTER,

            // BufferBoundsChecking
            BUFFER_SIZE,

            PAYLOAD_HEADER_SHORT,

            NUM_IMPLICIT_ARGS
        };

        /// @brief This set type is used to contain function arguments
        typedef std::set<int> ArgValSet;

        /// @brief This map type is used to map implicit argument types to
        ///        sets that contain the relevant Values.
        typedef std::map<ImplicitArg::ArgType, ArgValSet> ArgMap;

        typedef std::pair<ImplicitArg::ArgType, unsigned int> StructArgElement;
        typedef std::vector<StructArgElement> StructArgList;

        /// @brief  LLVM Type
        enum ValType {
            BYTE,
            SHORT,
            INT,
            LONG,
            FP32,
            CONSTPTR,
            PRIVATEPTR,
            GLOBALPTR
        };

        enum ValAlign {
            ALIGN_QWORD,
            ALIGN_DWORD,
            ALIGN_GRF,
            ALIGN_PTR
        };

        ImplicitArg(
            const ArgType&                  argType,
            const std::string&              name,
            const ValType                   valType,
            WIAnalysis::WIDependancy        dependency,
            unsigned int                    nbElement,
            ValAlign                        align,
            bool                            isConstantBuf);

        ImplicitArg(
            const ArgType& argType,
            const std::string& name,
            const ValType                   valType,
            WIAnalysis::WIDependancy        dependency,
            unsigned int                    nbElement,
            ValAlign                        align,
            bool                            isConstantBuf,
            llvm::GenISAIntrinsic::ID       GenIntrinsicID);

        /// @brief  Getter functions
        ArgType                         getArgType() const;
        const std::string&              getName() const;
        llvm::Type*                     getLLVMType(llvm::LLVMContext& context) const;
        WIAnalysis::WIDependancy        getDependency() const;
        unsigned int                    getAllocateSize(const llvm::DataLayout& DL) const;
        unsigned int                    getNumberElements() const;
        VISA_Type                       getVISAType(const llvm::DataLayout& DL) const;
        IGC::e_alignment                getAlignType(const llvm::DataLayout& DL) const;
        size_t                          getAlignment(const llvm::DataLayout& DL) const;
        unsigned int                    getPointerSize(const llvm::DataLayout& DL) const;
        bool                            isConstantBuf() const;
        bool                            isLocalIDs() const;
        llvm::GenISAIntrinsic::ID       getGenIntrinsicID() const;

    private:
        const ArgType                   m_argType;
        const std::string               m_name;
        const ValType                   m_valType;
        const WIAnalysis::WIDependancy  m_dependency;
        const unsigned int              m_nbElement;
        const ValAlign                  m_align;
        const bool                      m_isConstantBuf;
        const llvm::GenISAIntrinsic::ID m_GenIntrinsicID;
    };

    /// @brief  ImplicitArgs is used for accessing the actual implict information that is passed from
    ///         the OpenCL runtime to IGC.
    /// @Author Marina Yatsina
    class ImplicitArgs {
    public:
        // Constructors
        ImplicitArgs() {}

        /// @brief  Constructor.
        ///         Constructs the function's implicit arguments based on the given function's metadata
        ///         It actually constructs a mapping to a subset of IMPLICIT_ARGS
        /// @param  func the function to get impilcit args for.
        /// @param  the metadata utils object
        ImplicitArgs(const llvm::Function& func, const IGCMD::MetaDataUtils* pMdUtils);

        /// @brief  Returns the number of implicit arguments that are passed from the runtime
        /// @return The number of implicit arguments
        unsigned int size() const;

        /// @brief  Returns the i-th implicit argument
        /// @param  i               The implicit argument index
        /// @return The i-th implicit argument
        const ImplicitArg& operator[](unsigned int i) const;

        /// @brief  Returns the index of the appropriate implicit argument based on the given argument type
        /// @param  argType         The implicit argument type
        /// @return The implicit argument's index for a given argument type
        unsigned int getArgIndex(ImplicitArg::ArgType argType) const;

        /// @brief  Returns the index of the appropriate implicit image or sampler argument
        ///         based on the given argument type and the associated image argument
        /// @param  argType         The implicit argument type
        ///                         (height/width/depth for images, normalized/address/snapwa for samplers)
        /// @param  image           The associated image/sampler argument
        /// @return The implicit argument's index for a given argument type
        unsigned int getImageArgIndex(ImplicitArg::ArgType argType, const llvm::Argument* image) const;

        /// @brief  Returns the index of the appropriate implicit argument
        ///         based on the given argument type and the argument number
        /// @param  argType         The implicit argument type
        ///                         (height/width/depth for images, normalized/address/snapwa for samplers)
        /// @param  argNum          The explicit argument number of the type
        /// @return The implicit argument's index for a given argument type
        unsigned int getNumberedArgIndex(ImplicitArg::ArgType argType, int argNum) const;

        /// @brief  Returns if an implicit arg exists for the given explicit argument
        ///         based on the argument type and argument number
        /// @param  argType         The implicit argument type
        ///                         (height/width/depth for images, normalized/address/snapwa for samplers)
        /// @param  argNum          The explicit argument number of the type
        bool isImplicitArgExistForNumberedArg(ImplicitArg::ArgType argType, int argNum) const;

        /// @brief  Returns the argument type of the argument at the given index
        /// @param  i               The implicit argument index
        /// @return The argument type of the argument at the given index
        ImplicitArg::ArgType getArgType(unsigned int i) const;

        /// @brief  Returns the argument type of the given matching intrinsic ID
        /// @param  i               The GenISAIntrinsic ID
        /// @return The argument type
        static ImplicitArg::ArgType getArgType(llvm::GenISAIntrinsic::ID id);

        /// @brief  Returns the argument dependency of the given matching intrinsic ID
        /// @param  i               The GenISAIntrinsic ID
        /// @return The argument dependency
        static IGC::WIAnalysis::WIDependancy getArgDep(llvm::GenISAIntrinsic::ID id);

        /// @brief  Returns if the given arg type supports the GenISAIntrinsic instruction
        /// @param  i               The ImplicitArg type
        /// @return If intrinsic supported
        static bool hasIntrinsicSupport(ImplicitArg::ArgType i);

        /// @brief  Returns the explicit argument number of the given implicit argument index
        /// @param  i               The implicit argument index
        /// @return The explicit argument number of the given implicit argument index
        int32_t getExplicitArgNum(uint index) const;

        /// @brief  Returns the structure offset of the given implicit argument index
        /// @param  i               The implicit argument index
        /// @return The structure offset of the given implicit argument index
        int32_t getStructArgOffset(uint index) const;

        /// @brief  Creates implict arguments metadata for the given function based on the given implicit arguments
        ///         it should receive. If implicit metadata exists, it adds to it.
        /// @param  F               The function for which to create the implicit argument's metadata
        /// @param  implicitArgs    The implicit argument that are required by the given function
        /// @param  pMdUtils The Metadata API object
        static void addImplicitArgs(llvm::Function& F, const llvm::SmallVectorImpl<ImplicitArg::ArgType>& implicitArgs, const IGCMD::MetaDataUtils* pMdUtils);
        //TODO doc

        /// @brief  Creates implict arguments metadata for the given function amd all callers
        ///         based on the given implicit arguments it should receive. If implicit metadata exists, it adds to it.
        /// @param  F               The function for which to create the implicit argument's metadata
        /// @param  implicitArgs    The implicit argument that are required by the given function
        /// @param  pMdUtils The Metadata API object
        static void addImplicitArgsTotally(llvm::Function& F, const llvm::SmallVectorImpl<ImplicitArg::ArgType>& implicitArgs, const IGCMD::MetaDataUtils* pMdUtils);

        /// @brief  Creates implict image arguments metadata for the given function based on the given implicit image
        ///         arguments it should receive. If implicit image metadata exists, it adds to it.
        /// @param  F The function for which to create the implicit argument's metadata
        /// @param  argMap          A map of implict argument types to the Value pointers to the arguments
        /// @param  pMdUtils The Metadata API object
        static void addImageArgs(llvm::Function& F, const ImplicitArg::ArgMap& argMap, const IGCMD::MetaDataUtils* pMdUtils);
        // TODO doc

        static void addStructArgs(llvm::Function& F, const llvm::Argument* A, const ImplicitArg::StructArgList& S, const IGCMD::MetaDataUtils* pMdUtils);

        /// @brief  Creates implict arguments metadata for the given function based on the given implicit arguments
        ///         it should receive. If implicit metadata exists, it adds to it.
        /// @param  F               The function for which to create the implicit argument's metadata
        /// @param  argMap          A map of implict argument types to the set of numbers of arguments
        /// @param  pMdUtils        The Metadata API object
        static void addNumberedArgs(llvm::Function& F, const ImplicitArg::ArgMap& argMap, const IGCMD::MetaDataUtils* pMdUtils);

        /// @brief  Create implicit arguments metadata for the given function. It adds one
        ///         implicit argument for each explicit pointer argument to global or constant buffer.
        /// @param  F               The function for which to create the implicit argument's metadata
        /// @param  pMdUtils        The Metadata API object
        static void addBufferOffsetArgs(llvm::Function& F, const IGCMD::MetaDataUtils* pMdUtils, IGC::ModuleMetaData* modMD);

        /// @brief  Create implicit arguments metadata for the given function. It adds one
        ///         implicit argument for each explicit pointer argument to global or constant buffer.
        /// @param  F               The function for which to create the implicit argument's metadata
        /// @param  pMdUtils        The Metadata API object
        static void addBindlessOffsetArgs(llvm::Function& F, const IGCMD::MetaDataUtils* pMdUtils, IGC::ModuleMetaData* modMD);

        /// @brief  Check if the given implicit argument type exist in the(implicit) function argument associated
        /// @param  argType         The type of the implict argument that should be checked
        /// @return true if the argument exist, false otherwise.
        bool isImplicitArgExist(ImplicitArg::ArgType argType) const;
        static bool isImplicitArgExist(llvm::Function& F, ImplicitArg::ArgType argType, const IGCMD::MetaDataUtils* pMdUtils);

        /// @brief  Returns the (implicit) function argument associated with the given implicit argument type
        /// @param  F           The Function for which the implict argument should be returned
        /// @param  argType         The type of the implict argument that should be returned
        /// @return The (implicit) function argument associated with the given argument type
        ///         In case the argument doesn't exist, return nullptr
        llvm::Argument* getImplicitArg(llvm::Function &F, ImplicitArg::ArgType argType) const;

        /// @brief  Returns the (implicit) function argument associated with the given implicit argument type
        ///         and argument number
        /// @param  F        The Function for which the implict argument should be returned
        /// @param  argType  The type of the implict argument that should be returned
        /// @param  argNum   The explicit argument number of the type
        /// @return The (implicit) function argument associated with the given argument type and number
        ///         In case the argument doesn't exist, return nullptr
        llvm::Argument* getNumberedImplicitArg(llvm::Function &F, ImplicitArg::ArgType argType, int argNum) const;

        /// @brief  Returns true if the given argument type is an image or sampler.
        /// @param  argType The argument type to check.
        static bool isImplicitImage(ImplicitArg::ArgType argType);

        /// @brief  Returns true if the given argument type is a struct
        /// @param  argType The argument type to check.
        static bool isImplicitStruct(ImplicitArg::ArgType argType);

        /// @brief  Returns true if the given argument is an implicit argument.
        /// @param  arg The argument to check.
        bool isImplicitArg(llvm::Argument *arg) const;

        /// @brief  Returns explicit argument number associated with the given implicit argument.
        /// @param  implicitArg The associated implicit argument.
        int getExplicitArgNumForArg(llvm::Argument *implicitArg) const;

        llvm::Value* getImplicitArgValue(llvm::Function& F, ImplicitArg::ArgType argType, const IGCMD::MetaDataUtils* pMdUtils);

    private:
        /// @brief The function's metadata information.
        IGCMD::FunctionInfoMetaDataHandle m_funcInfoMD;

        static bool isImplicitArgExist(
            const IGCMD::FunctionInfoMetaDataHandle& funcInfo,
            ImplicitArg::ArgType argType);
    };

} // namespace IGC
