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

Copyright (C) 2017-2023 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "llvm/Config/llvm-config.h"
#include "AdaptorCommon/RayTracing/RayTracingAddressSpaceAliasAnalysis.h"
#include "Compiler/Optimizer/OpenCLPasses/AddressSpaceAliasAnalysis/AddressSpaceAliasAnalysis.h"
#include "Compiler/CodeGenPublic.h"
#include "Compiler/IGCPassSupport.h"
#include "Probe/Assertion.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvmWrapper/Analysis/AliasAnalysis.h"
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;
using namespace IGC;

namespace {
    class AddressSpaceAAResult : public IGCLLVM::AAResultBaseWrapper<AddressSpaceAAResult> {
        using BaseT = IGCLLVM::AAResultBaseWrapper<AddressSpaceAAResult>;
        friend BaseT;
        const TargetLibraryInfo& TLI;
        const CodeGenContext& CGC;
    public:
        explicit AddressSpaceAAResult(
            const TargetLibraryInfo& TLI,
            const CodeGenContext& ctx)
            : BaseT(), TLI(TLI), CGC(ctx) {}
        AddressSpaceAAResult(AddressSpaceAAResult&& Arg)
            : BaseT(std::move(Arg)), TLI(Arg.TLI), CGC(Arg.CGC) {}

        AddressSpaceAAResult(const AddressSpaceAAResult&) = delete;
        AddressSpaceAAResult& operator=(const AddressSpaceAAResult&) = delete;
        AddressSpaceAAResult& operator=(AddressSpaceAAResult&&) = delete;

        IGCLLVM::AliasResultEnum alias(const MemoryLocation& LocA, const MemoryLocation& LocB,
            AAQueryInfo & AAQI,
#if LLVM_VERSION_MAJOR < 16
            const llvm::Instruction* CtxI = nullptr
#else
            const llvm::Instruction* CtxI
#endif
        )
        {
            // DO NOT strip any casting as the address space is encoded in pointer
            // type. For `addrspacecast`, the current implementation in LLVM is too
            // aggressive to strip them. With address space resolution, we should not
            // have trivial cases where a non-generic pointer is cased into a generic
            // one.

            PointerType* PtrTy1 = dyn_cast<PointerType>(LocA.Ptr->getType());
            PointerType* PtrTy2 = dyn_cast<PointerType>(LocB.Ptr->getType());

            if (!PtrTy1 || !PtrTy2)
                return IGCLLVM::AliasResultEnum::NoAlias;

            unsigned AS1 = PtrTy1->getAddressSpace();
            unsigned AS2 = PtrTy2->getAddressSpace();

            // If LocA and LocB are pointers to different non-generic address spaces
            // then they do not alias.  This is not true in general, however, in OpenCL
            // all address spaces except generic are disjoint.
            //
            // Address spaces greater than ADDRESS_SPACE_NUM_ADDRESSES are used for
            // stateful accesses and may alias.
            //
            if (AS1 < ADDRESS_SPACE_NUM_ADDRESSES &&
                AS2 < ADDRESS_SPACE_NUM_ADDRESSES &&
                AS1 != ADDRESS_SPACE_GENERIC &&
                AS2 != ADDRESS_SPACE_GENERIC &&
                AS1 != AS2)
            {
                bool isDisjointMemory = true;
                if (CGC.allocatePrivateAsGlobalBuffer() && CGC.getModuleMetaData()->genericAccessesResolved)
                {
                    const bool isPrivateGlobal =
                        (AS1 == ADDRESS_SPACE_PRIVATE && AS2 == ADDRESS_SPACE_GLOBAL) ||
                        (AS1 == ADDRESS_SPACE_GLOBAL && AS2 == ADDRESS_SPACE_PRIVATE);

                    isDisjointMemory = !isPrivateGlobal;
                }

                if(isDisjointMemory)
                    return IGCLLVM::AliasResultEnum::NoAlias;
            }


            // Shared local memory doesn't alias any stateful memory.
            if ((AS1 == ADDRESS_SPACE_LOCAL && AS2 > ADDRESS_SPACE_NUM_ADDRESSES) ||
                (AS1 > ADDRESS_SPACE_NUM_ADDRESSES && AS2 == ADDRESS_SPACE_LOCAL))
            {
                return IGCLLVM::AliasResultEnum::NoAlias;
            }

            // Private memory doesn't alias any stateful memory
            if ((AS1 == ADDRESS_SPACE_PRIVATE && AS2 > ADDRESS_SPACE_NUM_ADDRESSES) ||
                (AS1 > ADDRESS_SPACE_NUM_ADDRESSES && AS2 == ADDRESS_SPACE_PRIVATE))
            {
                return IGCLLVM::AliasResultEnum::NoAlias;
            }


            /// For some client APIs (e.g. vulkan) compiler is free to assume that
            /// resources bound to two different bindings points never alias unless a
            /// resource is explicitly marked as being aliased.
            IGC_ASSERT(CGC.getModuleMetaData());
            if (AS1 > ADDRESS_SPACE_NUM_ADDRESSES &&
                AS2 > ADDRESS_SPACE_NUM_ADDRESSES &&
                CGC.getModuleMetaData()->statefulResourcesNotAliased)
            {
                bool isDirectAccess[2] = { false, false };
                unsigned resourceIndex[2] = { 0, 0 };
                BufferType resourceType[2] = { BUFFER_TYPE_UNKNOWN, BUFFER_TYPE_UNKNOWN };

                resourceType[0] = DecodeAS4GFXResource(AS1, isDirectAccess[0], resourceIndex[0]);
                resourceType[1] = DecodeAS4GFXResource(AS2, isDirectAccess[1], resourceIndex[1]);

                /// Returns true if resource type is a texture, constant buffer or a storage buffer.
                auto IsBufferOrTexture = [](
                    BufferType type)
                {
                    return (type == BufferType::CONSTANT_BUFFER || type == BufferType::UAV || type == BufferType::RESOURCE);
                };

                if (IsBufferOrTexture(resourceType[0]) &&
                    IsBufferOrTexture(resourceType[1]))
                {
                    if ((resourceType[0] != resourceType[1]) || // different resource types
                        (isDirectAccess[0] && isDirectAccess[1] && resourceIndex[0] != resourceIndex[1])) // direct access to different BTIs
                    {
                        return IGCLLVM::AliasResultEnum::NoAlias;
                    }
                }
            }

            return BaseT::alias(LocA, LocB, AAQI, CtxI);
        }

        bool pointsToConstantMemory(const llvm::MemoryLocation& Loc,
            AAQueryInfo & AAQI,
            bool OrLocal) {
            // Pointers to constant address space memory, well, point to constant memory
            PointerType* ptrType = dyn_cast<PointerType>(Loc.Ptr->getType());
            if (ptrType) {
                if (ptrType->getAddressSpace() == ADDRESS_SPACE_CONSTANT)
                    return true;

                bool DirectIdx;
                unsigned BufId;
                BufferType BufTy = DecodeAS4GFXResource(ptrType->getAddressSpace(),
                    DirectIdx, BufId);
                if (BufTy == CONSTANT_BUFFER)
                    return true;
            }

            return BaseT::pointsToConstantMemory(Loc, AAQI, OrLocal);
        }
    };

    class AddressSpaceAAWrapperPass : public ImmutablePass {
        std::unique_ptr<AddressSpaceAAResult> Result;

    public:
        static char ID;

        AddressSpaceAAWrapperPass();

        void getAnalysisUsage(AnalysisUsage& AU) const override {
            AU.setPreservesAll();
            AU.addRequired<TargetLibraryInfoWrapperPass>();
            AU.addRequired<CodeGenContextWrapper>();
        }

        bool doInitialization(Module& M) override {
            if(M.size() > 0)
            {
#if LLVM_VERSION_MAJOR >= 10
                auto& F = *M.begin(); // see llvmWrapper/Analysis/TargetLibraryInfo.h
#endif
                Result.reset(new AddressSpaceAAResult(
                    getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(),
                    *getAnalysis<CodeGenContextWrapper>().getCodeGenContext()));
            }
            return false;
        }

        bool doFinalization(Module& M) override {
            Result.reset();
            return false;
        }

        AddressSpaceAAResult& getResult() { return *Result; }
        const AddressSpaceAAResult& getResult() const { return *Result; }
    };

} // End anonymous namespace

#define PASS_FLAG     "igc-address-space-alias-analysis"
#define PASS_DESC     "Address space alias analysis"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS true
IGC_INITIALIZE_PASS_BEGIN(AddressSpaceAAWrapperPass, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
IGC_INITIALIZE_PASS_END(AddressSpaceAAWrapperPass, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY, PASS_ANALYSIS)

char AddressSpaceAAWrapperPass::ID = 0;

AddressSpaceAAWrapperPass::AddressSpaceAAWrapperPass() : ImmutablePass(ID) {
    initializeAddressSpaceAAWrapperPassPass(*PassRegistry::getPassRegistry());
}

ImmutablePass* IGC::createAddressSpaceAAWrapperPass() {
    return new AddressSpaceAAWrapperPass();
}

void IGC::addJointAddressSpaceAAResults(Pass& P, Function& F, AAResults& AAR) {
    if (auto * WrapperPass = P.getAnalysisIfAvailable<AddressSpaceAAWrapperPass>())
        AAR.addAAResult(WrapperPass->getResult());
    addRayTracingAddressSpaceAAResult(P, F, AAR);
}

// Wrapper around LLVM's ExternalAAWrapperPass so that the default constructor
// gets the callback.
class IGCExternalAAWrapper : public ExternalAAWrapperPass {
public:
  static char ID;

  IGCExternalAAWrapper() : ExternalAAWrapperPass(&addJointAddressSpaceAAResults) {}
};

ImmutablePass* IGC::createIGCExternalAAWrapper() {
    return new IGCExternalAAWrapper();
}

char IGCExternalAAWrapper::ID = 0;

IGC_INITIALIZE_PASS_BEGIN(IGCExternalAAWrapper, "igc-aa-wrapper", "IGC Address space alias analysis Wrapper", PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_END(IGCExternalAAWrapper, "igc-aa-wrapper", "IGC Address space alias analysis Wrapper", PASS_CFG_ONLY, PASS_ANALYSIS)
