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

Copyright (C) 2020-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

//===----------------------------------------------------------------------===//
///
/// Provides a utility to scan up a few instructions looking for a base address
/// corresponding to a region in raytracing memory.  This is to be reused
/// in whatever passes need info on what memory regions they are touching.
///
//===----------------------------------------------------------------------===//

#include "Compiler/CodeGenPublic.h"
#include "MemRegionAnalysis.h"
#include "MDFrameWork.h"
#include "Probe/Assertion.h"
#include "API/RayDispatchGlobalData.h"
#include "llvmWrapper/IR/Value.h"
#include <optional>

using namespace llvm;
using namespace IGC;

static bool isGlobalPtr(const InlineDataIntrinsic *I, const IGC::ModuleMetaData& moduleMetaData)
{
    uint32_t globalPtrOffset = 0;

    if (moduleMetaData.rtInfo.GlobalDataStyle == RayDispatchInlinedDataStyle::Xe)
    {
        constexpr uint32_t GlobalPtrOffset =
            offsetof(RayDispatchInlinedData, RayDispatchGlobalDataPtr) / sizeof(uint64_t);
        globalPtrOffset = GlobalPtrOffset;
    }

    return (I->getArg() == globalPtrOffset);
}

static std::optional<RTMemRegion> getIntrinsicRegion(
    const GenIntrinsicInst* GII,
    const IGC::ModuleMetaData& moduleMetaData)
{
    switch (GII->getIntrinsicID())
    {
    case GenISAIntrinsic::GenISA_SWHotZonePtr:
        return RTMemRegion::SWHotZone;
    case GenISAIntrinsic::GenISA_GlobalBufferPointer:
        return RTMemRegion::RTGlobals;
    case GenISAIntrinsic::GenISA_LocalBufferPointer:
        return RTMemRegion::LocalArgs;
    case GenISAIntrinsic::GenISA_AsyncStackPtr:
        return RTMemRegion::RTAsyncStack;
    case GenISAIntrinsic::GenISA_SyncStackPtr:
        return RTMemRegion::RTSyncStack;
    case GenISAIntrinsic::GenISA_SWStackPtr:
    case GenISAIntrinsic::GenISA_ContinuationSignpost:
        return RTMemRegion::SWStack;
    case GenISAIntrinsic::GenISA_InlinedData:
        if (isGlobalPtr(cast<InlineDataIntrinsic>(GII), moduleMetaData))
            return RTMemRegion::RTGlobals;
        else
            return std::nullopt;
    default:
        return std::nullopt;
    }
}

namespace IGC {

std::optional<RTMemRegion>
getRTRegionByAddrspace(const Value* V, const ModuleMetaData &MMD)
{
    auto* PtrTy = dyn_cast<PointerType>(V->getType());
    if (!PtrTy)
        return std::nullopt;

    uint32_t Addrspace = PtrTy->getPointerAddressSpace();

    auto& rtInfo = MMD.rtInfo;

    if (Addrspace == rtInfo.RTAsyncStackAddrspace)
        return RTMemRegion::RTAsyncStack;
    else if (Addrspace == rtInfo.SWHotZoneAddrspace)
        return RTMemRegion::SWHotZone;
    else if (Addrspace == rtInfo.SWStackAddrspace)
        return RTMemRegion::SWStack;
    else if (Addrspace == rtInfo.RTSyncStackAddrspace)
        return RTMemRegion::RTSyncStack;

    return std::nullopt;
}

std::optional<RTMemRegion> getRegionOffset(
    const Value* Ptr,
    const IGC::ModuleMetaData& moduleMetaData,
    const DataLayout *DL,
    uint64_t* Offset,
    uint64_t* dereferenceable_value)
{
    // Set an initial value for the Offset, overwriting whatever garbage value there may be.
    // If there is a getelementptr instruction the value of the *Offset will be updated.
    if (Offset) {
        *Offset = 0;
    }

    while (Ptr)
    {
        if (auto* GEPI = dyn_cast<GetElementPtrInst>(Ptr))
        {
            Ptr = GEPI->getPointerOperand();

            if (Offset)
            {
                unsigned IdxWidth = DL->getIndexSizeInBits(
                    Ptr->getType()->getPointerAddressSpace());
                APInt BasePtrOffset(IdxWidth, 0);
                // If there is a constant offset on the getelementptr instruction like i64 4
                // we can evaluate it at compile time
                if (GEPI->accumulateConstantOffset(*DL, BasePtrOffset)) {
                    *Offset += BasePtrOffset.getZExtValue();
                }
                // If there is a non-constant or variable offset on the getelementptr instruction like i64 %x
                // we cannot evaluate it at compile time, only at run time
                else {
                    return std::nullopt;
                }
            }
        }
        else if (auto* BCI = dyn_cast<BitCastInst>(Ptr))
        {
            Ptr = BCI->getOperand(0);
        }
        else if (auto *GII = dyn_cast<GenIntrinsicInst>(Ptr))
        {
            if (dereferenceable_value) {
                bool CanBeNull = false;
                bool CanBeFreed = false;
                *dereferenceable_value = IGCLLVM::getPointerDereferenceableBytes(Ptr, *DL, CanBeNull, CanBeFreed);
                // We probably only want to set the *dereferenceable_value if CanBeNull is false
                // But CanBeNull should never be true if you're calling this function on any of those intrinsics in getRegionOffset()
                IGC_ASSERT(!CanBeNull);
            }

            return getIntrinsicRegion(GII, moduleMetaData);
        }
        else
        {
            Ptr = nullptr;
        }
    }

    return std::nullopt;
}

// This is conceptually the same idea as getRegionOffset() but it doesn't
// compute offsets and is intended to be used late in compilation after GEPs
// have been lowered (or at least mostly lowered).  It thus looks through
// more constructs that we don't care to look at in earlier stages.
//
// Currently, the only user of this is for late stage LSC cache controls
// determination.
std::optional<RTMemRegion> getRTRegion(const Value* V, const ModuleMetaData &MMD)
{
    if (auto Region = getRTRegionByAddrspace(V, MMD))
        return Region;

    while (V)
    {
        auto* I = dyn_cast<Instruction>(V);
        if (!I)
            return std::nullopt;

        switch (I->getOpcode())
        {
        case Instruction::GetElementPtr:
        case Instruction::BitCast:
        case Instruction::IntToPtr:
        case Instruction::PtrToInt:
        case Instruction::Add:
        case Instruction::Or:
            V = I->getOperand(0);
            break;
        case Instruction::Call:
            if (auto* GII = dyn_cast<GenIntrinsicInst>(I))
                return getIntrinsicRegion(GII, MMD);
            return std::nullopt;
        default:
            return std::nullopt;
        }
    }

    return std::nullopt;
}

} // namespace IGC

