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

Copyright (C) 2022-2024 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/PatternMatch.h>
#include <llvm/Pass.h>
#include <llvm/Support/Debug.h>
#include <llvm/Support/raw_ostream.h>
#include "common/LLVMWarningsPop.hpp"
#include "GenISAIntrinsics/GenIntrinsics.h"
#include "Compiler/CISACodeGen/ShaderCodeGen.hpp"
#include "Compiler/MetaDataUtilsWrapper.h"
#include "Compiler/CISACodeGen/CSWalkOrder.hpp"
#include "Probe/Assertion.h"

using namespace llvm;
using namespace IGC;


void IGC::selectWalkOrderInPass(
    bool useLinearWalk,
    uint numberOfTypedAccess,
    uint numberOfUntypedAccess,
    uint num1DAccesses,
    uint num2DAccesses,
    uint numSLMAccesses,
    uint threadGroupSize_X,
    uint threadGroupSize_Y,
    uint threadGroupSize_Z,
    CodeGenContext* ctx,
    SComputeShaderWalkOrder& walkOrderStruct)
{
    ModuleMetaData* MMD = ctx->getModuleMetaData();
    const CPlatform& platform = ctx->platform;
    const CDriverInfo& driverInfo = ctx->m_DriverInfo;
    ThreadIDLayout& threadIDLayout = walkOrderStruct.m_threadIDLayout;
    CS_WALK_ORDER& walkOrder = walkOrderStruct.m_walkOrder;
    EMIT_LOCAL_MASK &emitMask = walkOrderStruct.m_emitMask;
    bool& enableHWGenerateLID = walkOrderStruct.m_enableHWGenerateLID;

    if (MMD->csInfo.neededThreadIdLayout == ThreadIDLayout::QuadTile)
    {
        threadIDLayout = ThreadIDLayout::QuadTile;
        return;
    }

    bool is_pow2_x = iSTD::IsPowerOfTwo(threadGroupSize_X);
    bool is_pow2_y = iSTD::IsPowerOfTwo(threadGroupSize_Y);
    bool is_pow2_z = iSTD::IsPowerOfTwo(threadGroupSize_Z);
    if (IGC_IS_FLAG_ENABLED(SetDefaultTileYWalk) && is_pow2_x &&
        platform.enableSetDefaultTileYWalk() && driverInfo.SupportHWGenerateTID()) {
        threadIDLayout = ThreadIDLayout::TileY;
        walkOrder = CS_WALK_ORDER::WO_YXZ;
    }
    if ((numberOfTypedAccess >= numberOfUntypedAccess) &&
        threadGroupSize_Y % 4 == 0 &&
        !MMD->csInfo.disableLocalIdOrderOptimizations &&
        IGC_IS_FLAG_ENABLED(UseTiledCSThreadOrder)) {
        threadIDLayout = ThreadIDLayout::TileY;
        walkOrder = CS_WALK_ORDER::WO_YXZ;
    }

    bool needsLinearWalk =
        MMD->csInfo.neededThreadIdLayout == ThreadIDLayout::X;
    if (platform.supportHWGenerateTID() && driverInfo.SupportHWGenerateTID())
    {
        // If KeepTileYForFlattened == 2, use the platform default value.
        // Otherwise 0 is forced off, 1 is forced on.
        bool KeepTileYForFlattenedValue = IGC_GET_FLAG_VALUE(KeepTileYForFlattened) == 2 ?
            platform.EnableKeepTileYForFlattenedDefault() :
            IGC_IS_FLAG_ENABLED(KeepTileYForFlattened);
        if (!KeepTileYForFlattenedValue && useLinearWalk)
        {
            needsLinearWalk = true;
        }
    }
    if (needsLinearWalk)
    {
        threadIDLayout = ThreadIDLayout::X;
        walkOrder = CS_WALK_ORDER::WO_XYZ;
    }
    if (!(platform.supportHWGenerateTID() && driverInfo.SupportHWGenerateTID()))
        return;

    //if no LID is used ever, HWGenerateLID should be disabled
    if (EMIT_LOCAL_MASK::EM_NONE == emitMask) {
        enableHWGenerateLID = false;
        return;
    }

    //if TileY is selected AND EnableHWGenerateThreadIDForTileY==0, use sw to generate LID (legacy behavior)
    //ATM, we don't know for sure if tileY improves performance, so, we have a regkey here to tune it.
    //todo: this is for tuning only, remove this if logic once we know for sure if tileY should be enabled for XeHP+
    if ((threadIDLayout == ThreadIDLayout::TileY) &&
        !IGC_IS_FLAG_ENABLED(EnableHWGenerateThreadIDForTileY)) {
        enableHWGenerateLID = false;
        return;
    }

    //in case of HWGenerateLID, Y must be power2
    if ((threadIDLayout == ThreadIDLayout::TileY) &&
        !iSTD::IsPowerOfTwo(threadGroupSize_Y)) {
        threadIDLayout = ThreadIDLayout::X;
        walkOrder = CS_WALK_ORDER::WO_XYZ;
    }

    //if not all DIMs are used, we can assume not-used DIM vals are 1, so, HW might generate LIDs even if original not-used DIM val != pow2
    //say, (31, 17, 3), if only X dim is used, HW will generate LID for (31,1,1). If both XY are used, then, HW cannot generate LIDs
    if (EMIT_LOCAL_MASK::EM_X == emitMask) {
        threadGroupSize_Y = 1;
        threadGroupSize_Z = 1;
        //it makes no sense to use TileY if no Y is used at all. Disable it.
        threadIDLayout = ThreadIDLayout::X;
    }
    else if (EMIT_LOCAL_MASK::EM_XY == emitMask) {
        threadGroupSize_Z = 1;
    }//else if (EMIT_LOCAL_MASK::EM_XYZ == emitMask)

    // Will use linear for the case like 64x1x1.
    // Open: How about 1x32x1?
    if (threadIDLayout == ThreadIDLayout::TileY &&
        threadGroupSize_Y == 1 && threadGroupSize_Z == 1) {
        // 1D thread group
        threadIDLayout = ThreadIDLayout::X;
        walkOrder = CS_WALK_ORDER::WO_XYZ;
    }

    if (!IGC_IS_FLAG_ENABLED(EnableNonOCLWalkOrderSel) || needsLinearWalk) {
        if (threadGroupSize_Y == 1 && threadGroupSize_Z == 1)
        {
            walkOrder = CS_WALK_ORDER::WO_YZX;
            enableHWGenerateLID = true;
        }
        else if (threadGroupSize_X == 1 && threadGroupSize_Z == 1)
        {
            walkOrder = CS_WALK_ORDER::WO_XZY;
            enableHWGenerateLID = true;
        }
        else
        {
            walkOrder = CS_WALK_ORDER::WO_XYZ;
            enableHWGenerateLID = (is_pow2_x && is_pow2_y);
        }
        //disable tileY if walkorder cannot be changed
        threadIDLayout = ThreadIDLayout::X;
        overrideWalkOrderKeysInPass(is_pow2_x, is_pow2_y, is_pow2_z, walkOrderStruct, ctx);
        return;
    }

    // Perfer linear walk for 2D dispatch, but linear UAV surface
    if ((IGC_IS_FLAG_ENABLED(ForceLinearWalkOnLinearUAV) ||
        MMD->compOpt.ForceLinearWalkOnLinearUAV) &&
        (threadIDLayout == ThreadIDLayout::TileY) &&
        EMIT_LOCAL_MASK::EM_XY == emitMask &&
        num1DAccesses)
    {
        threadIDLayout = ThreadIDLayout::X;
        walkOrder = CS_WALK_ORDER::WO_XYZ;
    }

    // If EnableNewTileYCheck == 2, use the platform default value. Otherwise 0 is forced off, 1 is forced on.
    bool b_EnableNewTileYCheck = IGC_GET_FLAG_VALUE(EnableNewTileYCheck) == 2 ?
        platform.EnableNewTileYCheckDefault() : IGC_GET_FLAG_VALUE(EnableNewTileYCheck);
    if (b_EnableNewTileYCheck &&
        IGC_IS_FLAG_ENABLED(SetDefaultTileYWalk) &&
        (threadIDLayout == ThreadIDLayout::TileY) &&
        EMIT_LOCAL_MASK::EM_XY == emitMask)
    {
        // check 1D, 2D, SLM accesses
        int num1D = num1DAccesses + (int)(numSLMAccesses / 4);
        int num2D = num2DAccesses;
        if (num1D > num2D && num2D <= 5)
        {
            threadIDLayout = ThreadIDLayout::X;
            walkOrder = CS_WALK_ORDER::WO_XYZ;
        }
    }

    auto order = selectBestWalkOrderInPass(
        threadIDLayout, is_pow2_x, is_pow2_y, is_pow2_z);

    if (order) {
        walkOrder = *order;
        enableHWGenerateLID = true;
    }
    else {
        // Is 2D or 3D dispatch and isnt pow2, so the HW doesn't support it
        enableHWGenerateLID = false;
        threadIDLayout = ThreadIDLayout::X;
        walkOrder = CS_WALK_ORDER::WO_XYZ;
    }

    overrideWalkOrderKeysInPass(is_pow2_x, is_pow2_y, is_pow2_z, walkOrderStruct, ctx);
}

void IGC::overrideWalkOrderKeysInPass(
    bool is_pow2_x, bool is_pow2_y, bool is_pow2_z,
    SComputeShaderWalkOrder& walkOrderStruct,
    CodeGenContext* ctx)
{
    ThreadIDLayout& threadIDLayout = walkOrderStruct.m_threadIDLayout;
    CS_WALK_ORDER& walkOrder = walkOrderStruct.m_walkOrder;
    bool& enableHWGenerateLID = walkOrderStruct.m_enableHWGenerateLID;

    ModuleMetaData* MMD = ctx->getModuleMetaData();
    const CPlatform& platform = ctx->platform;
    const CDriverInfo& driverInfo = ctx->m_DriverInfo;

    if ((IGC_IS_FLAG_ENABLED(ForceTileY) || MMD->csInfo.forceTileYWalk) &&
        platform.supportHWGenerateTID() && driverInfo.SupportHWGenerateTID())
    {
        threadIDLayout = ThreadIDLayout::TileY;
        walkOrder = CS_WALK_ORDER::WO_YXZ;
        enableHWGenerateLID = enableHWGenerateLIDInPass(walkOrder, is_pow2_x, is_pow2_y, is_pow2_z);
    }

    if (MMD->csInfo.walkOrderEnabled)
    {
        walkOrder = (CS_WALK_ORDER)MMD->csInfo.walkOrderOverride;
        enableHWGenerateLID = enableHWGenerateLIDInPass(walkOrder, is_pow2_x, is_pow2_y, is_pow2_z);
    }

    if (IGC_IS_FLAG_ENABLED(OverrideCsWalkOrderEnable))
    {
        walkOrder = (CS_WALK_ORDER)IGC_GET_FLAG_VALUE(OverrideCsWalkOrder);
        enableHWGenerateLID = enableHWGenerateLIDInPass(walkOrder, is_pow2_x, is_pow2_y, is_pow2_z);
    }

    if (IGC_IS_FLAG_ENABLED(OverrideCsTileLayoutEnable))
    {
        threadIDLayout = (ThreadIDLayout)IGC_IS_FLAG_ENABLED(OverrideCsTileLayout);
    }
}

bool IGC::enableHWGenerateLIDInPass(
    CS_WALK_ORDER walk_order,
    bool is_pow2_x, bool is_pow2_y, bool is_pow2_z)
{
    bool bEnableHWGenerateLID = false;

    switch (walk_order)
    {
    case CS_WALK_ORDER::WO_XYZ:
    case CS_WALK_ORDER::WO_YXZ:
        bEnableHWGenerateLID = (is_pow2_x && is_pow2_y);
        break;

    case CS_WALK_ORDER::WO_XZY:
    case CS_WALK_ORDER::WO_ZXY:
        bEnableHWGenerateLID = (is_pow2_x && is_pow2_z);
        break;

    case CS_WALK_ORDER::WO_YZX:
    case CS_WALK_ORDER::WO_ZYX:
        bEnableHWGenerateLID = (is_pow2_y && is_pow2_z);
        break;
    }
    return bEnableHWGenerateLID;
}

Optional<CS_WALK_ORDER>
IGC::selectBestWalkOrderInPass(
    ThreadIDLayout Layout,
    bool is_pow2_x, bool is_pow2_y, bool is_pow2_z)
{
    constexpr uint UNDEF = std::numeric_limits<uint>::max();
    uint order0 = UNDEF;
    uint order1 = UNDEF;
    if (Layout == ThreadIDLayout::TileY)
    {
        IGC_ASSERT(is_pow2_y);
        order0 = 1;
        order1 = (is_pow2_x ? 0 : (is_pow2_z ? 2 : UNDEF));
    }
    else
    {
        //below is from HAS p-code except tileY
        //try to find walk_order so that HW can generate LID
        if (is_pow2_x)
        {
            // (pow2,pow2,z) or (pow2,y,pow2) or illegal
            order0 = 0;
            order1 = (is_pow2_y ? 1 : (is_pow2_z ? 2 : UNDEF));
        }
        else if (is_pow2_y)
        {
            // (x,pow2,pow2) or illegal
            order0 = 1;
            order1 = (is_pow2_z ? 2 : UNDEF);
        }
    }

    if (order1 != UNDEF)
    {
        // select walkorder
        return getWalkOrderInPass(order0, order1);
    }

    return None;
}

void IGC::setEmitLocalMaskInPass(SGVUsage channelNum, EMIT_LOCAL_MASK& emitMask)
{
    //only 4 patterns are supported: None; X; XY; XYZ
    switch (channelNum)
    {
    case THREAD_ID_IN_GROUP_X:
        emitMask = (EMIT_LOCAL_MASK::EM_NONE == emitMask) ? EMIT_LOCAL_MASK::EM_X : emitMask;
        break;
    case THREAD_ID_IN_GROUP_Y:
        emitMask = (EMIT_LOCAL_MASK::EM_NONE == emitMask || EMIT_LOCAL_MASK::EM_X == emitMask) ?
            EMIT_LOCAL_MASK::EM_XY : emitMask;
        break;
    case THREAD_ID_IN_GROUP_Z:
        emitMask = EMIT_LOCAL_MASK::EM_XYZ;
        break;
    default:
        break;
    }
}

//order0: the internal walk dim
//order1: the intermediate walk dim
//e.g.: 1, 0 means, YXZ walkorder
CS_WALK_ORDER IGC::getWalkOrderInPass(uint order0, uint order1)
{
    auto getWalkOrderValue = [](uint order0, uint order1) constexpr {
        return (order0 << 4 | order1 << 2);
        };

    switch (getWalkOrderValue(order0, order1))
    {
    case getWalkOrderValue(0, 1): return CS_WALK_ORDER::WO_XYZ; //012
    case getWalkOrderValue(0, 2): return CS_WALK_ORDER::WO_XZY; //021
    case getWalkOrderValue(1, 0): return CS_WALK_ORDER::WO_YXZ; //102
    case getWalkOrderValue(1, 2): return CS_WALK_ORDER::WO_YZX; //120
    case getWalkOrderValue(2, 0): return CS_WALK_ORDER::WO_ZXY; //201
    case getWalkOrderValue(2, 1): return CS_WALK_ORDER::WO_ZYX; //210
    default:
        IGC_ASSERT_MESSAGE(0, "unhandled case!");
        return CS_WALK_ORDER::WO_XYZ;
    }
}