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

Copyright (C) 2023 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "Compiler/Optimizer/OpenCLPasses/GEPLoopStrengthReduction/GEPLoopStrengthReduction.hpp"
#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include <llvm/Analysis/LoopInfo.h>
#include <llvm/Analysis/ScalarEvolution.h>
#include <llvm/Analysis/ScalarEvolutionExpressions.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/GetElementPtrTypeIterator.h>
#include "llvm/Support/Debug.h"
#include <llvm/Transforms/Utils/Local.h>
#include "llvmWrapper/Transforms/Utils/LoopUtils.h"
#include "llvmWrapper/Transforms/Utils/ScalarEvolutionExpander.h"
#include "llvmWrapper/IR/Instructions.h"
#include "common/LLVMWarningsPop.hpp"

#include "Probe/Assertion.h"
#include "common/igc_regkeys.hpp"
#include "common/IGCIRBuilder.h"
#include "Compiler/CISACodeGen/IGCLivenessAnalysis.h"
#include "Compiler/CodeGenPublic.h"

using namespace llvm;
using namespace IGC;

#define DEBUG_TYPE "gep-loop-strength-reduction"

// This pass reduces strength of address calculations inside loops. Xe architecture doesn't have native
// support for access from "base pointer plus offset". When access "pointer[offset]" is compiled into GEP
// instruction "getelementptr %constantPointer, %variableOffset", in codegen compiler must translate it
// into "(long long)pointer + offset*sizeof(*pointer)", generating additional mov/add/shl instructions.
// If access happens in the loop, it might be beneficial to replace it into "getelementptr %variablePointer,
// %constantOffset". This generates less instruction for address calculation in codegen.
//
//
// This pass does three kinds of optimizations:
//
// 1. If GEP's index is constant between loop iterations, move GEP to loop preheader. This is very similar
//    to LICM, with one difference: if loop contains multiple GEPs with the same pointer and different
//    offsets, LICM moves all GEPs to preheader, but this pass moves only base GEP and keeps remaining
//    address calculations inside loops, but replace them with constant offsets. This gives less register
//    pressure than LICM (only one address in preheader) and still simplifies address calculation in
//    loop body.
//
// 2. If GEP's index in incremented in constant steps in the loop, move GEP for base address calculation
//    to loop preheader and add it as new induction variable to the loop. This pass translates code:
//
//      int id = get_global_id(0) * n_iters;
//      for (int i = 0; i < n_iters; ++i) {
//        buf[id + i] = buf[id + i] + 3.0f;
//      }
//
//    Into:
//
//      int id = get_global_id(0) * n_iters;
//      global float *p = buf + id;
//      for (int i = 0; i < n_iters; ++i, ++p) {
//        *p = *p + 3.0f;
//      }
//
//    GEP instruction with multiple indices are also supported if last index is the only variable. In this
//    case new induction variable is direct pointer to indexed type. Before reduction:
//
//      for (int i = 0; i < n_iters; ++i) {
//        s->buf[id + i] = 3.0f; // s = struct { int float[] }
//      }
//
//    After:
//
//      global float *p = s->buf + id;
//      for (int i = 0; i < n_iters; ++i, ++p) {
//        *p = 3.0f;
//      }
//
// 3. If GEP's index in incremented in constant steps in the loop, but base address calculation can't be
//    moved to loop preheader, simplify index calculation for other GEPs using the same index plus constant
//    offset. This will translate code:
//
//      int id = get_global_id(0) * n_iters;
//      for (int i = 0; i < n_iters; i+=2) {
//        buf[id + i + 1] = buf[id + i] + 3.0f;
//      }
//
//    Into:
//
//      int id = get_global_id(0) * n_iters;
//      for (int i = 0; i < n_iters; i+=2) {
//        global float *p = buf + id + i;
//        *(p + 1) = *p + 3.0f;
//      }


enum ReductionType
{
    // Moves GEP for base address calculation to loop's preheader. If address changes inside loop
    // with constant step, adds pointer as new induction variable to the loop. All GEPs using the
    // same base address + constant offset are modified to use new pointer.
    REDUCE_TO_PREHEADER,

    // Does not move GEP to loop's preheader and does not add new induction variable to the
    // loop. Base address for first GEP is recalculated on each iteration, but all related
    // GEPs are modified to use the same pointer + constant offset, reducing address calculation
    // inside loop. Doesn't increase register pressure.
    REDUCE_INDEX_ONLY,
};


// Score used by heuristic for deciding what type of reduction to apply.
struct Score
{
    Score() : ReducesInstructions(false), RegisterPressure(0), ContainsMuli64(false) {}

    // True if reduction to preheader would lower number of instructions in
    // loop. False otherwise.
    bool ReducesInstructions;

    // Estimated increase in register pressure when reducing to loop's preheader.
    unsigned RegisterPressure;

    // Flag to show presence of i64 mul instructions in the address calculation.
    bool ContainsMuli64;
};


// Single GEP instruction that is a candiate for reduction. It is expressed by SCEV:
//     { start, +, step }
// Where step is constant integer (including 0), and start can be calculated in loop's preheader.
struct ReductionCandidate
{
    ReductionCandidate(GetElementPtrInst *GEP, const SCEV *S, int64_t Delta)
        : GEP(GEP), S(S), Delta(Delta) {}

    bool operator==(const ReductionCandidate &RHS) const
    {
        return GEP == RHS.GEP && S == RHS.S && Delta == RHS.Delta;
    }

    bool isValidForReduction(const Loop *L, const DominatorTree *DT);

    bool isBetterForReduction(const ReductionCandidate &Other);

    void print(raw_ostream &OS);

    // Instruction to reduce
    GetElementPtrInst *GEP;

    // SCEV expressing starting index
    const SCEV *S;

    // Index delta between group's base pointer and this pointer
    int64_t Delta;
};


// Sometimes reduction candidates might be grouped. This is possible if both candidates:
// 1) delta between start values can be expressed in constant int
// 2) have the same step increment value (including 0)
//
// Example: This loop has three uses of "output" buffer with three different offsets.
//
//    int id = get_global_id(0);
//    for (int i = 32; i < n_iters - 32; i += 32) {
//        output[id + i] = output[id + i + 32] * output[id + i - 32];
//    }
//
// Each access could be reduced independently, but since all three offsets are increased in
// the same constant step, it is more beneficial to reduce them to one common pointer:
//
//    int id = get_global_id(0);
//    global float *outputm32 = output + id;
//    for (int i = 32; i < n_iters - 32; i += 32) {
//        *(outputm32 + 32) = *(outputm32 + 64) * *outputm32;
//        outputm32 += 32;
//    }
class ReductionCandidateGroup
{
friend class Scorer;

public:

    ReductionCandidateGroup(Loop *L, const DominatorTree *DT, GetElementPtrInst *GEP, const SCEV *S, const SCEV *Step)
        : L(L), DT(DT), Step(Step), Base(GEP, S, 0), RT(REDUCE_TO_PREHEADER) {}

    bool addToGroup(ScalarEvolution &SE, GetElementPtrInst *GEP, const SCEV *S, const SCEV *Step);

    void transform(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E);

    bool isValid();

    // Return how many GEP instructions will be reduced.
    uint32_t getGEPCount() { return 1 + Others.size(); }

    // Returns reduced GEPs. Filled after successful reduction.
    const SmallVectorImpl<GetElementPtrInst*> &getReduced() { return Reduced; }

    ReductionType getReductionType() { return RT; }
    void setReductionType(ReductionType RT) { this->RT = RT; }

    Loop *getLoop() { return L; }

    Score getScore() { return Score; }

    void print(raw_ostream &OS);

private:

    void append(GetElementPtrInst *GEP, const SCEV *S, int64_t Delta);
    void append(ReductionCandidate &C) { append(C.GEP, C.S, C.Delta); }

    ReductionCandidate &getCheapestCandidate();

    void swapBase(ReductionCandidate &C);

    void reduceToPreheader(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E);
    void reduceIndexOnly(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E);

    Value *getStepValue(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E, BasicBlock *BB);

    // Base GEP to reduce
    ReductionCandidate Base;

    // Other GEPs that can be reduced together with base
    SmallVector<ReductionCandidate, 4> Others;

    // Increment step value
    const SCEV *Step;

    // Group member with the smallest SCEV expression. Preferred for reduction
    // to preheader, as it should give smallest increase in register pressure.
    std::optional<ReductionCandidate> Cheapest;

    // All GEPs that were reduced and are safe to delete.
    SmallVector<GetElementPtrInst*, 8> Reduced;

    Loop *L;

    const DominatorTree *DT;

    Score Score;

    ReductionType RT;
};


// Scores reduction candidates:
//   1. Estimates removed instructions by reducing "base_ptr + index * sizeof(type)".
//   2. Estimates increase in register pressure.
class Scorer
{
public:
    Scorer(const DataLayout &DL, ModuleMetaData &MMD, IGCLivenessAnalysis &RPE, WIAnalysisRunner &WI)
        : DL(DL), MMD(MMD), RPE(RPE), WI(WI) {}

    void score(SmallVectorImpl<ReductionCandidateGroup> &Candidates);

private:
    void score(ReductionCandidateGroup &Candidate);

    void scoreReducedInstructions(ReductionCandidateGroup &Candidate);
    void scoreRegisterPressure(ReductionCandidateGroup &Candidate);

    int estimateIndexInstructions(const Loop &L, GetElementPtrInst *GEP, bool &ContainsMuli64);
    int estimatePointerAddition(ReductionCandidateGroup &Candidate);

    const DataLayout &DL;
    IGCLivenessAnalysis &RPE;
    ModuleMetaData &MMD;
    WIAnalysisRunner &WI;
};


// Analyzes GEP instructions in a single loop and selects candidates for reduction.
class Analyzer
{
public:
    Analyzer(const DataLayout &DL, DominatorTree &DT, Loop &L, LoopInfo &LI, ScalarEvolution &SE, SCEVExpander &E)
        : DT(DT), L(L), LI(LI), SE(SE), E(E) {}

    void analyze(SmallVectorImpl<ReductionCandidateGroup> &Result);

private:

    // Represents deconstructed SCEV expression { start, +, step }.
    // Start SCEV will be used to calculate base pointer, and Step SCEV
    // will increase new induction variable on each iteration.
    struct DeconstructedSCEV
    {
        DeconstructedSCEV() : Start(nullptr), Step(nullptr), ConvertedMulExpr(false) {}

        bool isValid();

        const SCEV* Start;
        const SCEV* Step;

        // True if input SCEV:
        //   x * { start, +, step }
        // Was converted into:
        //   { x * start, +, x * step }
        bool ConvertedMulExpr;
    };

    void analyzeGEP(GetElementPtrInst *GEP);
    bool doInitialValidation(GetElementPtrInst *GEP);

    bool deconstructSCEV(const SCEV *S, DeconstructedSCEV &Result);

    DominatorTree &DT;
    Loop &L;
    LoopInfo &LI;
    ScalarEvolution &SE;
    SCEVExpander &E;

    SmallVector<ReductionCandidateGroup, 16> Candidates;
};


// Tracks estimated register pressure.
class RegisterPressureTracker
{
public:

    RegisterPressureTracker(Function &F, CodeGenContext &CGC, IGCLivenessAnalysis &RPE, IGCFunctionExternalRegPressureAnalysis &FRPE, WIAnalysisRunner &WI);

    bool fitsPressureThreshold(ReductionCandidateGroup &C);
    void updatePressure(ReductionCandidateGroup &C, SCEVExpander &E);

    void trackDeletedInstruction(Value *V);

private:

    unsigned MaxAllowedPressure;
    unsigned FunctionExternalPressure;

    IGCLivenessAnalysis &RPE;
    WIAnalysisRunner &WI;

    // Basic Blocks impacted by reduction, requiring register pressure reestimation.
    BBSet BBsToUpdate;

    // Keep track what new instructions inserted by SCEV Expander were already added to estimation.
    SmallPtrSet<Instruction*, 32> VisitedNewInsts;
};


// Does reduction on collected candidates.
class Reducer
{
public:
    Reducer(IGCLLVM::IRBuilder<> &IRB, RegisterPressureTracker &RPT, SCEVExpander &E, bool AllowLICM)
        : IRB(IRB), RPT(RPT), E(E), AllowLICM(AllowLICM) {}

    bool reduce(SmallVectorImpl<ReductionCandidateGroup> &Candidates);

private:

    void cleanup(ReductionCandidateGroup &C);

    IGCLLVM::IRBuilder<> &IRB;
    RegisterPressureTracker &RPT;
    SCEVExpander &E;

    bool AllowLICM;
};


// Set of functions/classes helping manipulating SCEV objects.
namespace SCEVHelper
{
    const SCEV *dropExt(const SCEV *S);

    bool isValid(const SCEV *S);
    bool isEqual(const SCEV *A, const SCEV *B);

    // ScalarEvolution::getAddExpr requires all operands to have the same
    // type. Extend type if required.

    // Builds SCEVAddExpr instance. Function ScalarEvolution::getAddExpr requires all
    // operands to have the same type. This class wraps ScalarEvolution::getAddExpr,
    // but extends operands if it is needed to keep them all in one type.
    class SCEVAddBuilder
    {
    public:
        SCEVAddBuilder(ScalarEvolution &SE) : SE(SE) {}

        SCEVAddBuilder &add(const SCEV *S, bool Negative = false);

        SCEVAddBuilder &addNegative(const SCEV *S) { return add(S, true); }

        const SCEV *build();

    private:

        struct Op
        {
            Op(const SCEV *S, bool Negative) : S(S), Negative(Negative) {}

            const SCEV *S;
            bool Negative;
        };

        ScalarEvolution &SE;
        SmallVector<Op, 16> Ops;
    };

    // Builds SCEVMulExpr instance. Function ScalarEvolution::getMulExpr requires all
    // operands to have the same type. This class wraps ScalarEvolution::getMulExpr,
    // but extends operands if it is needed to keep them all in one type.
    class SCEVMulBuilder
    {
    public:
        SCEVMulBuilder(ScalarEvolution &SE) : SE(SE) {}

        SCEVMulBuilder &add(const SCEV *S);

        const SCEV *build();

    private:
        ScalarEvolution &SE;
        SmallVector<const SCEV*, 4> Ops;
    };
};


class GEPLoopStrengthReduction : public llvm::FunctionPass
{
public:
    static char ID; // Pass identification, replacement for typeid

    GEPLoopStrengthReduction(bool AllowLICM = true);

    ~GEPLoopStrengthReduction() {}

    virtual llvm::StringRef getPassName() const override
    {
        return "GEPLoopStrengthReduction";
    }

    void getAnalysisUsage(llvm::AnalysisUsage &AU) const override;

    virtual bool runOnFunction(llvm::Function &F) override;

private:
    bool AllowLICM;
};


void ReductionCandidate::print(raw_ostream &OS)
{
    OS << "{gep=";
    GEP->printAsOperand(OS, false);
    OS << ", ptr=";
    GEP->getPointerOperand()->printAsOperand(OS, false);
    OS << ", scev=";
    S->print(OS);
    OS << ", delta=";
    OS << Delta;
    OS << "}";
}


void ReductionCandidateGroup::print(raw_ostream &OS)
{
    OS << "{loop=";
    OS << L->getName();
    OS << ", geps=" << Others.size() + 1;
    OS << ", base=";
    Base.print(OS);
    OS << ", others=[";
    for (auto it = Others.begin(); it != Others.end();)
    {
        it->print(OS);
        ++it;
        if (it != Others.end())
            OS << ", ";
    }
    OS << "], step=";
    Step->print(OS);
    OS << "}";
}


// Considering loop:
//
//   for (i ...)
//   {
//     if (cond)
//     {
//       x[i - 1] = ...
//     }
//     x[i] = ...
//   }
//
// If address is calculated in if-case, if there is no guarantee it would be calculated on the
// first iteration, it can't be reduced to preheader, as value could be invalid.
bool ReductionCandidate::isValidForReduction(const Loop* L, const DominatorTree* DT)
{
    if (!DT->dominates(GEP->getParent(), L->getLoopLatch()))
        return false;

    // Even if address is always calculated, it might not always be used.
    // Find at least one always-executed use to proof address is safe.
    for (auto U = GEP->users().begin(); U != GEP->users().end(); ++U)
    {
        if (auto *I = dyn_cast<Instruction>(*U))
        {
            if (L->contains(I) && DT->dominates(I->getParent(), L->getLoopLatch()))
                return true;
        }
    }

    return false;
}


bool ReductionCandidate::isBetterForReduction(const ReductionCandidate &Other)
{
    // If lowest delta is selected as base, remaining GEPs in the group will use
    // positive indices only. This should make other load/store optimizations easier.
    return (Delta < Other.Delta) ||
           (Delta == Other.Delta && S->getExpressionSize() < Other.S->getExpressionSize());
}


// If possible, adds another reduction candidate to this group. Reductions
// can be grouped together if:
//   1) delta between start values can be expressed in constant int
//   2) have the same increment step value
// New candidate can take position of group's base if it uses less instructions
// to calculate.
// Returns true if candidate was added to group.
bool ReductionCandidateGroup::addToGroup(ScalarEvolution &SE, GetElementPtrInst *GEP, const SCEV *S, const SCEV *Step)
{
    if (!SCEVHelper::isEqual(this->Step, Step))
        return false;

    if (Base.GEP->getPointerOperand() != GEP->getPointerOperand())
        return false;

    if (Base.GEP->getType() != GEP->getType())
        return false;

    // Compare indices (except last one)
    if (Base.GEP->getNumIndices() != GEP->getNumIndices())
        return false;

    if (!std::equal(Base.GEP->indices().begin(), Base.GEP->indices().end() - 1, GEP->indices().begin()))
        return false;

    // Can't use ScalarEvolution::computeConstantDifference, as it only
    // supports SCEVAddExpr with two operands. Calculate difference as:
    //     new candidate's operands + (-1 * base's operands)
    SCEVHelper::SCEVAddBuilder Builder(SE);
    const SCEVConstant *Sum = dyn_cast<SCEVConstant>(Builder.add(S).addNegative(Base.S).build());
    if (!Sum)
        return false;

    int64_t Delta = Sum->getValue()->getSExtValue();

    // Delta is constant int, reductions can be grouped together.
    append(GEP, S, Delta);

    return true;
}


void ReductionCandidateGroup::append(GetElementPtrInst *GEP, const SCEV *S, int64_t Delta)
{
    Others.emplace_back(GEP, S, Delta);

    // Group changed, forget which member is the cheapest.
    Cheapest.reset();
}


// Sets group's base GEP to provided candidate.
void ReductionCandidateGroup::swapBase(ReductionCandidate &C)
{
    if (Base == C)
        return;

    auto *ToReplace = std::find(Others.begin(), Others.end(), C);
    IGC_ASSERT(ToReplace != Others.end());

    ReductionCandidate NewBase = *ToReplace;

    // Directly replace old slot with old base.
    *ToReplace = Base;

    // Update candidates with new delta.
    for (auto &O : Others)
        O.Delta -= NewBase.Delta;

    Base = NewBase;
    Base.Delta = 0;
}


ReductionCandidate &ReductionCandidateGroup::getCheapestCandidate()
{
    if (Cheapest)
        return Cheapest.value();

    ReductionCandidate *BestCandidate = nullptr;

    for (auto It = Others.begin(); It != Others.end(); ++It)
    {
        if (!It->isValidForReduction(L, DT))
            continue;

        if (!BestCandidate || It->isBetterForReduction(*BestCandidate))
            BestCandidate = It;
    }

    if (BestCandidate && (!Base.isValidForReduction(L, DT) || BestCandidate->isBetterForReduction(Base)))
        return Cheapest.emplace(*BestCandidate);

    return Cheapest.emplace(Base);
}


bool ReductionCandidateGroup::isValid()
{
    if (Base.isValidForReduction(L, DT))
        return true;

    for (auto It = Others.begin(); It != Others.end(); ++It)
    {
        if (It->isValidForReduction(L, DT))
            return true;
    }

    return false;
}


void ReductionCandidateGroup::transform(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E)
{
    switch (RT)
    {
    case REDUCE_TO_PREHEADER:
        return reduceToPreheader(IRB, E);
    case REDUCE_INDEX_ONLY:
    default:
        return reduceIndexOnly(IRB, E);
    }
}


// Transforms all grouped GEP instructions:
//   1) Calculate start value in loop's preheader
//   2) Add this new pointer as loop's induction variable with constant step
//   3) Replace GEP instructions to use this pointer + constant offset
void ReductionCandidateGroup::reduceToPreheader(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E)
{
    // Updates group's base to candidate with smallest SCEV expression.
    swapBase(getCheapestCandidate());

    // Create pointer's starting value in preheader
    BasicBlock *LPH = L->getLoopPreheader();
    const SCEV *Start = Base.S;
    Value *StartIndex = E.expandCodeFor(Start, Start->getType(), &LPH->back());

    IRB.SetInsertPoint(&LPH->back());
    SmallVector<Value*, 4> Indices(Base.GEP->indices().begin(), Base.GEP->indices().end() - 1);
    Indices.push_back(StartIndex);
    Value *Pointer = IRB.CreateGEP(Base.GEP->getSourceElementType(), Base.GEP->getPointerOperand(), Indices);
    Type* ptrElTy = IGCLLVM::getGEPIndexedType(Base.GEP->getSourceElementType(), Indices);

    // Create phi node if pointer is moved in loop
    if (!Step->isZero())
    {
        // Add new phi node with pointer as induction variable
        SmallVector<BasicBlock*, 4> Latches;
        L->getLoopLatches(Latches);

        IRB.SetInsertPoint(L->getHeader(), L->getHeader()->begin());
        PHINode *Phi = IRB.CreatePHI(Pointer->getType(), Latches.size() + 1);

        Phi->addIncoming(Pointer, LPH);

        // In each latch increment pointer and add to phi node
        for (auto *L : Latches)
        {
            IRB.SetInsertPoint(&L->back());
            Value *Inc = IRB.CreateGEP(ptrElTy, Phi, getStepValue(IRB, E, L));
            Phi->addIncoming(Inc, L);
        }

        Pointer = Phi;
    }

    // Replace base GEP
    Base.GEP->replaceAllUsesWith(Pointer);
    Reduced.push_back(Base.GEP);

    // Replace remaining GEPs in the group
    for (auto &Other : Others)
    {
        IRB.SetInsertPoint(Other.GEP);
        Other.GEP->replaceAllUsesWith(Other.Delta == 0 ? Pointer : IRB.CreateGEP(ptrElTy, Pointer, IRB.getInt64(Other.Delta)));
        Reduced.push_back(Other.GEP);
    }
}


// Keeps base GEP untouched and transforms other GEPs in the group to use
// base GEP + constant offset.
void ReductionCandidateGroup::reduceIndexOnly(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E)
{
    for (auto &Other : Others)
    {
        // TODO If GEPs instructions are in different BBs, base GEP could be moved
        // to one common position (e.g. at the start of the iteration) to allow
        // other GEPs to share the same pointer. But it increases pressure inside loop?
        if (Base.GEP->getParent() == Other.GEP->getParent())
        {
            IRB.SetInsertPoint(Other.GEP);
            Other.GEP->replaceAllUsesWith(Other.Delta == 0 ? Base.GEP :
                IRB.CreateGEP(Base.GEP->getResultElementType(), Base.GEP, IRB.getInt64(Other.Delta)));
            Reduced.push_back(Other.GEP);
        }
    }
}


Value *ReductionCandidateGroup::getStepValue(IGCLLVM::IRBuilder<> &IRB, SCEVExpander &E, BasicBlock *BB)
{
    if (auto *S = dyn_cast<SCEVConstant>(Step))
        return IRB.getInt64(dyn_cast<SCEVConstant>(Step)->getValue()->getSExtValue());

    if (auto *S = dyn_cast<SCEVUnknown>(Step))
        return S->getValue();

    return E.expandCodeFor(Step, Step->getType(), &BB->back());
}


void Scorer::score(SmallVectorImpl<ReductionCandidateGroup> &Candidates)
{
    for (auto &C : Candidates)
    {
        score(C);
    }
}


void Scorer::score(ReductionCandidateGroup &Candidate)
{
    scoreReducedInstructions(Candidate);
    scoreRegisterPressure(Candidate);
}


// Estimate if reduction will remove instructions from loop. Takes into account instructions
// added when lowering GEP to direct pointer arithmetics.
void Scorer::scoreReducedInstructions(ReductionCandidateGroup &Candidate)
{
    auto &Cheapest = Candidate.getCheapestCandidate();

    // Score "index * sizeof(type) + base_ptr". Split into parts:
    //   1. "index"          - number of instructions required to calculate value (might be 0
    //                         if index is just loop's induction variable).
    //   2. "* sizeof(type)" - single "shl" instruction
    //   3. "+ base_ptr"     - single "add" instruction
    int score = 0;

    if (!Candidate.Step->isZero())
    {
        // Reduction adds new instruction - incrementation of new induction variable at the end
        // of the iteration.
        score -= 1;
    }

    // Score "* sizeof(type)"
    if (unsigned(DL.getTypeSizeInBits(Cheapest.GEP->getResultElementType())) > 8)
        score += 1;

    // Score "+ base_ptr"
    score += estimatePointerAddition(Candidate);
    bool ContainsMuli64 = false;

    // Score "index"
    score += estimateIndexInstructions(*Candidate.getLoop(), Cheapest.GEP, ContainsMuli64);

    Candidate.Score.ReducesInstructions = score > 0;
    Candidate.Score.ContainsMuli64 = ContainsMuli64;
}


// Estimate if compiler will emit "base_ptr + offset" instruction.
int Scorer::estimatePointerAddition(ReductionCandidateGroup &Candidate)
{
    if (Candidate.Base.GEP->getPointerAddressSpace() != ADDRESS_SPACE_LOCAL)
        return 1;

    // Local address space pointers are replaced to constant integers by GenIRLowering
    // pass (CG phase), giving more chances to optimize pointer arithmetics.

    FunctionMetaData &FMD = MMD.FuncMD[Candidate.Base.GEP->getFunction()];

    // Check if base pointer is zero ("0 + offset" will be optimized out).
    for (auto &Offets : FMD.localOffsets)
    {
        if (Candidate.Base.GEP->getPointerOperand() == Offets.m_Var)
        {
            if (Offets.m_Offset == 0)
                return 0;
            break; // exit early
        }
    }

    return 1;
}


// Estimates how many instructions required to calculate index would be reduced to preheader.
// This differs from checking SCEV expression size, which it might represent simplified index
// calculation.
// Sets ContainsMuli64 flag to show if i64 multiplication is present in the gep index calculation.
int Scorer::estimateIndexInstructions(const Loop &L, GetElementPtrInst *GEP, bool &ContainsMuli64)
{
    Instruction *Index = dyn_cast<Instruction>(*(GEP->operands().end() - 1));
    if (!Index)
        return 0;

    // Keep track of visited instructions (don't go into cycles).
    SmallPtrSet<Instruction*, 8> Visited;

    std::function<int(Instruction*)> Visit = [&](Instruction *I)
    {
        if (Visited.insert(I).second == false)
            return 0;

        if (!L.contains(I))
            return 0;

        int instructions = 0;

        switch (I->getOpcode())
        {
        case Instruction::PHI:
            if (I->getParent() == L.getHeader())
                return 0; // no need to go further
            break;
        case Instruction::SExt:
        case Instruction::ZExt:
            break;
        default:
            instructions += 1;
        }

        if (I->getOpcode() == Instruction::Mul && I->getType()->isIntegerTy(64))
            ContainsMuli64 = true;

        for (auto *It = I->operands().begin(); It != I->operands().end(); ++It)
        {
            if (auto *Next = dyn_cast<Instruction>(It))
                instructions += Visit(Next);
        }

        return instructions;
    };

    return Visit(Index);
}


// Returns estimated increase in pressure.
void Scorer::scoreRegisterPressure(ReductionCandidateGroup &Candidate)
{
    auto &Cheapest = Candidate.getCheapestCandidate();

    // Method SCEVExpander::isHighCostExpansion could be a good candidate to use here, but:
    //   1. It doesn't return increase value, only boolean if expansion fits in given budget.
    //   2. Depends on correctly defined TargetTransformInfo.

    // To estimate increase in register pressure, iterate over all instructions inside loop
    // that used to calculate index and estimate register size. In reality, change in
    // pressure might be different, as SCEVExpander might simplify calculations.

    Instruction* Index = dyn_cast<Instruction>(*(Cheapest.GEP->operands().end() - 1));
    if (!Index)
        return;

    auto *L = Candidate.getLoop();
    auto *F = Cheapest.GEP->getParent()->getParent();
    uint SIMD = numLanes(RPE.bestGuessSIMDSize(F));

    ValueSet Instructions;

    // Keep track of visited instructions (don't go into cycles).
    ValueSet Visited;

    std::function<void(Instruction*)> Visit = [&](Instruction* I)
    {
        if (Visited.insert(I).second == false)
            return;

        if (!L->contains(I))
            return;

        switch (I->getOpcode())
        {
        case Instruction::PHI:
            if (I->getParent() == L->getHeader())
                return; // no need to go further
            break;
        default:
            Instructions.insert(I);
        }

        for (auto* It = I->operands().begin(); It != I->operands().end(); ++It)
        {
            if (auto* Next = dyn_cast<Instruction>(It))
                Visit(Next);
        }
    };

    Visit(Index);

    uint SizeInBytes = RPE.estimateSizeInBytes(Instructions, *F, SIMD, &WI);

    Candidate.Score.RegisterPressure = RPE.bytesToRegisters(SizeInBytes);
}


void Analyzer::analyze(SmallVectorImpl<ReductionCandidateGroup> &Result)
{
    // Require simple loop with one entry and one latch.
    if (!IGCLLVM::isInnermost(&L) || !L.isLoopSimplifyForm() || !L.getLoopLatch())
        return;

    Candidates.clear();

    for (auto *BB : L.getBlocks()) {
        for (auto &I : *BB) {
            if (auto *GEP = dyn_cast<GetElementPtrInst>(&I))
                analyzeGEP(GEP);
        }
    }

    // After grouping candidates together, it's possible some groups don't have
    // at least one valid pointer to use as a base for reduction.
    Candidates.erase(
        std::remove_if(Candidates.begin(), Candidates.end(),
            [](ReductionCandidateGroup& C) {
                if (C.isValid())
                    return false;
                LLVM_DEBUG(
                    dbgs() << "  dropping candidate with no GEP valid for reduction: ";
                    C.print(dbgs());
                    dbgs() << "\n");
                return true;
            }),
        Candidates.end());

    Result.append(Candidates.begin(), Candidates.end());
}


// Checks if GEP instruction can be transformed.
void Analyzer::analyzeGEP(GetElementPtrInst *GEP)
{
    if (!doInitialValidation(GEP))
        return;

    Value *Index = *(GEP->indices().end() - 1);

    const SCEV *S = SE.getSCEV(Index);
    if (!SCEVHelper::isValid(S))
        return;

    Analyzer::DeconstructedSCEV Result;
    if (!deconstructSCEV(S, Result))
        return;

    if (!Result.isValid())
        return;

    const SCEV* Start = Result.Start;
    const SCEV* Step = Result.Step;

    if (S->getType() != Start->getType())
        Start = isa<SCEVSignExtendExpr>(S) ? SE.getSignExtendExpr(Start, S->getType()) : SE.getZeroExtendExpr(Start, S->getType());

    // Try adding reduction to existing group
    for (auto &c : Candidates)
    {
        if (c.addToGroup(SE, GEP, Start, Step))
        {
            LLVM_DEBUG(
                dbgs() << "  updated group of candidates: ";
                c.print(dbgs());
                dbgs() << "\n");
            return;
        }
    }

    Candidates.emplace_back(&L, &DT, GEP, Start, Step);
    LLVM_DEBUG(
        dbgs() << "  found new candidate to optimize: ";
        Candidates.back().print(dbgs());
        dbgs() << "\n");
}

// Does early GEP instruction analysis, before reaching for SCEV analysis.
// Returns false if GEP can be dropped early.
bool Analyzer::doInitialValidation(GetElementPtrInst *GEP)
{
    // Optimize only explicit global/local memory access. Private memory might be
    // promoted to GRFs, keep them as is.
    auto AddressSpace = GEP->getPointerAddressSpace();
    if (AddressSpace != ADDRESS_SPACE_GLOBAL && AddressSpace != ADDRESS_SPACE_LOCAL && AddressSpace != ADDRESS_SPACE_GENERIC)
        return false;

    // If pointer is instruction, it must be usable in loop preheader.
    if (auto *I = dyn_cast<Instruction>(GEP->getPointerOperand()))
    {
        if (!DT.dominates(I, L.getLoopPreheader()) && I->getParent() != L.getLoopPreheader())
            return false;
    }

    // TODO: Pointer could be inttopointer cast of value available in loop preheader.
    // In this case transformation is still possible if casting is moved to preheader.

    // Only last index can be variable
    if (!std::all_of(GEP->indices().begin(), GEP->indices().end() - 1, [](Value *V) { return isa<Constant>(V); }))
        return false;
    Value *Index = *(GEP->indices().end() - 1);

    // Make sure last indexed type is array. Don't try to mix up and optimize access
    // to different fields in aggregate type.
    if (GEP->getNumIndices() > 1)
    {
        SmallVector<Value*, 8> Indices(GEP->indices().begin(), GEP->indices().end() - 1);
        Type *Ty = GetElementPtrInst::getIndexedType(GEP->getSourceElementType(), Indices);
        if (!Ty || !Ty->isArrayTy())
            return false;
    }

    // Don't reduce if index is used outside of loop to access the same pointer.
    // TODO: These accesses could be modified to also use pointer induction variable
    // added by this pass.
    SmallPtrSet<Value*, 8> Visited = { GEP }; // keep track of visited instructions
    std::function<bool(Value*)> CheckOutsideAccess = [&](Value *V)
    {
        if (Visited.insert(V).second == false)
            return false;

        for (auto It = V->users().begin(); It != V->users().end(); ++It)
        {
            if (auto *Phi = dyn_cast<PHINode>(*It))
            {
                if (CheckOutsideAccess(Phi))
                    return true;
            }

            if (auto *OtherGEP = dyn_cast<GetElementPtrInst>(*It))
            {
                if (!L.contains(OtherGEP))
                    return true;
            }
        }

        return false;
    };
    if (CheckOutsideAccess(Index))
        return false;

    return true;
}

// Takes SCEV expression returned by ScalarEvolution and deconstructs it into
// expected format { start, +, step }. Returns false if expressions can't be
// parsed and reduced.
bool Analyzer::deconstructSCEV(const SCEV *S, Analyzer::DeconstructedSCEV &Result)
{
    // Drop ext instructions to analyze nested content.
    S = SCEVHelper::dropExt(S);

    // First check if expression can be fully expanded in preheader. If so, no need
    // to process is further, but instead treat expression as:
    //   { start, +, 0 }
    // This will do LICM-like reduction moving GEP to preheader, without adding new
    // induction variable.
    if (SE.isLoopInvariant(S, &L))
    {
        Result.Start = S;
        Result.Step = SE.getConstant(Type::getInt64Ty(L.getHeader()->getContext()), 0);
        return true;
    }

    // Expect SCEV expression:
    //   { start, +, step }
    // where step is constant
    if (auto *Add = dyn_cast<SCEVAddRecExpr>(S))
    {
        // Consider the following SCEV as a mul pattern:
        // {(%val * %multiplier),+,%multiplier}
        auto *Op0 = Add->getOperand(0);
        auto *Op1 = Add->getOperand(1);
        bool Conv = false;
        if (auto *Mul = dyn_cast<SCEVMulExpr>(Op0)) {
            auto *MulOp0 = Mul->getOperand(0);
            auto *MulOp1 = Mul->getOperand(1);
            if (MulOp0 == Op1 || MulOp1 == Op1) {
                Conv = true;
            }
        }

        if (!Add->isAffine())
            return false;

        if (Add->getNumOperands() != 2)
            return false;

        const SCEV *OpStep = Add->getOperand(1);

        // Step must be constant in loop's body.
        if (!SE.isLoopInvariant(OpStep, &L))
            return false;

        Result.Start = Add->getStart();
        Result.Step = OpStep;
        Result.ConvertedMulExpr = Conv;

        return IGCLLVM::isSafeToExpandAt(Result.Start, &L.getLoopPreheader()->back(), &SE, &E);
    }

    // If expression is:
    //   x + { start, +, step }
    // then change it to:
    //   { start + x, +, step }
    //
    // It is possible that "x" is not constant inside loop, but is recalculated
    // on every iteration. In this case it is not a valid scenario for reduction
    // and will be dropped by IGCLLVM::isSafeToExpandAt.
    if (auto *Add = dyn_cast<SCEVAddExpr>(S))
    {
        // There can be only one expression with step != 0.
        Result.Step = SE.getConstant(Type::getInt64Ty(L.getHeader()->getContext()), 0);

        SCEVHelper::SCEVAddBuilder Builder(SE);

        for (auto *Op : Add->operands())
        {
            Analyzer::DeconstructedSCEV OpResult;

            if (!deconstructSCEV(Op, OpResult))
                return false;

            if (!OpResult.Step->isZero())
            {
                if (!Result.Step->isZero())
                    return false; // unsupported expression with multiple steps
                Result.Step = OpResult.Step;
            }

            Builder.add(OpResult.Start);
        }

        Result.Start = Builder.build();

        return IGCLLVM::isSafeToExpandAt(Result.Start, &L.getLoopPreheader()->back(), &SE, &E);
    }

    // If expression is:
    //   x * { start, +, step }
    // then change it to:
    //   { x * start, +, x * step }
    //
    // Warning: GEP's new index will not be a constant integer, but a new SCEV expression.
    if (auto *Mul = dyn_cast<SCEVMulExpr>(S))
    {
        // SCEVAddRecExpr will be SCEV with step != 0. Any other SCEV is a multiplier.
        bool FoundAddRec = false;
        SCEVHelper::SCEVMulBuilder StartBuilder(SE), StepBuilder(SE);

        for (auto *Op : Mul->operands())
        {
            Analyzer::DeconstructedSCEV OpResult;
            if (!deconstructSCEV(Op, OpResult))
                return false;

            if (OpResult.Step->isZero())
            {
                StartBuilder.add(OpResult.Start);
                StepBuilder.add(OpResult.Start);
            }
            else
            {
                if (FoundAddRec)
                    return false; // unsupported expression with multiple SCEVAddRecExpr
                FoundAddRec = true;

                StartBuilder.add(OpResult.Start);
                StepBuilder.add(OpResult.Step);
            }
        }

        if (!FoundAddRec)
            return false;

        Result.Start = StartBuilder.build();
        Result.Step = StepBuilder.build();
        Result.ConvertedMulExpr = true;

        if (!SE.isLoopInvariant(Result.Step, &L))
            return false;

        return IGCLLVM::isSafeToExpandAt(Result.Start, &L.getLoopPreheader()->back(), &SE, &E);
    }

    return false;
}


bool Analyzer::DeconstructedSCEV::isValid()
{
    if (!Start || !Step)
        return false;

    // Validate step.
    auto Ty = SCEVHelper::dropExt(Step)->getSCEVType();

    if (Ty == scConstant)
        return true;

    bool IsMul = Ty == scMulExpr || ConvertedMulExpr;
    if (IsMul && IGC_IS_FLAG_ENABLED(EnableGEPLSRMulExpr))
        return true;

    return IGC_IS_FLAG_ENABLED(EnableGEPLSRUnknownConstantStep);
}


RegisterPressureTracker::RegisterPressureTracker(Function &F, CodeGenContext &CGC, IGCLivenessAnalysis &RPE, IGCFunctionExternalRegPressureAnalysis &FRPE, WIAnalysisRunner &WI)
    : RPE(RPE), WI(WI)
{
    MaxAllowedPressure = static_cast<unsigned>(CGC.getNumGRFPerThread() * IGC_GET_FLAG_VALUE(GEPLSRThresholdRatio) / 100.0f);

    FunctionExternalPressure = FRPE.getExternalPressureForFunction(&F);
}


void RegisterPressureTracker::trackDeletedInstruction(Value *V)
{
    if (auto *I = dyn_cast<Instruction>(V))
        BBsToUpdate.insert(I->getParent());

    for (auto It = RPE.getInSet().begin(); It != RPE.getInSet().end(); ++It)
    {
        if (It->second.count(V))
            BBsToUpdate.insert(It->first);
    }

    for (auto It = RPE.getOutSet().begin(); It != RPE.getOutSet().end(); ++It)
    {
        if (It->second.count(V))
            BBsToUpdate.insert(It->first);
    }
}


bool RegisterPressureTracker::fitsPressureThreshold(ReductionCandidateGroup &C)
{
    BasicBlock *Preheader = C.getLoop()->getLoopPreheader();
    auto *F = Preheader->getParent();
    uint SIMD = numLanes(RPE.bestGuessSIMDSize(F));

    unsigned MaxLoopPressure = RPE.getMaxRegCountForLoop(*C.getLoop(), SIMD, &WI);
    unsigned AdditionalPressure = C.getScore().RegisterPressure;

    InsideBlockPressureMap BBListing;
    RPE.collectPressureForBB(*Preheader, BBListing, SIMD, &WI);
    unsigned LoopExternalPressureInBytes = BBListing[cast<Value>(Preheader->getTerminator())];
    unsigned LoopExternalPressure = RPE.bytesToRegisters(LoopExternalPressureInBytes);

    unsigned InitialPressure = FunctionExternalPressure + MaxLoopPressure;
    unsigned EstimatedPressure = InitialPressure + AdditionalPressure;

    // Try not to increase register pressure above threshold.
    if (EstimatedPressure >= MaxAllowedPressure)
    {
        // Even if the optimization icnreases register pressure, apply it in case we can move mul i64 to preheader.
        // This heuristic is based on the fact that mul i64 is expensive instruction and potential spills are generated out of the loop.
        unsigned NewInternalLoopPressure = LoopExternalPressure - MaxLoopPressure + AdditionalPressure;
        if (C.getScore().ContainsMuli64 && NewInternalLoopPressure < MaxAllowedPressure) {
            return true;
        }

        LLVM_DEBUG(
            dbgs() << "  Estimated register pressure " << EstimatedPressure << " above threshold " << MaxAllowedPressure << "; can't fully reduce ";
            C.print(dbgs());
            dbgs() << "\n");
        return false;
    }

    return true;
}


void RegisterPressureTracker::updatePressure(ReductionCandidateGroup &C, SCEVExpander &E)
{

#if LLVM_VERSION_MAJOR < 14
    BBsToUpdate.clear();
    Function *F = C.getLoop()->getLoopPreheader()->getParent();
    RPE.rerunLivenessAnalysis(*F);
    return;
#else

    // Refresh all BBs in loop.
    BBsToUpdate.insert(C.getLoop()->getBlocks().begin(), C.getLoop()->getBlocks().end());

    if (C.getReductionType() == REDUCE_TO_PREHEADER)
    {
        // When adding new induction variable to loop, we need to refresh loop's preheader and body (all BBs).
        // SCEVExpander can push new instructions to outer loops, including adding new induction variables.
        // Find outermost loop touched by SCEVExpander and refresh register estimation for it.

        // Query SCEVExpander for new instructions.
        auto AllInsertedInstructions = E.getAllInsertedInstructions();

        // Start searching from inner loop.
        Loop *TopLoop = C.getLoop();

        for (auto *I : AllInsertedInstructions)
        {
            // Skip if instruction was inserted by previous reductions.
            if (VisitedNewInsts.insert(I).second == false)
                continue;

            if (TopLoop->contains(I))
                continue;

            if (TopLoop->getLoopPreheader() == I->getParent())
                continue;

            // At this point we know that instruction is not in current loop. Check outer loops.

            Loop *L;
            for (L = TopLoop->getParentLoop(); L != nullptr; L = L->getParentLoop())
            {
                if (L->contains(I) || L->getLoopPreheader() == I->getParent())
                    break;
            }

            if (L)
            {
                // Found new loop.
                TopLoop = L;
            }
            else
            {
                LLVM_DEBUG(
                    dbgs() << "  Found instruction in basic block [" << I->getParent()->getName() << "] not present in any loop; instruction: ";
                    I->print(dbgs(), true);
                    dbgs() << "\n");
                BBsToUpdate.insert(I->getParent());
            }
        }

        BBsToUpdate.insert(TopLoop->block_begin(), TopLoop->block_end());
        BBsToUpdate.insert(TopLoop->getLoopPreheader());
    }

    // BBs with removed instructions should be already collected in Reducer::cleanup.

    Function *F = C.getLoop()->getLoopPreheader()->getParent();

    if (!BBsToUpdate.empty())
    {
        RPE.rerunLivenessAnalysis(*F, &BBsToUpdate);
        BBsToUpdate.clear();
    }
#endif // LLVM_VERSION_MAJOR
}


bool Reducer::reduce(SmallVectorImpl<ReductionCandidateGroup> &Candidates)
{
    if (Candidates.empty())
        return false;

    bool changed = false;

    // If LICM is disabled, drop all candidates that would move only one GEP instruction to preheader.
    if (!AllowLICM)
    {
        Candidates.erase(
            std::remove_if(Candidates.begin(), Candidates.end(),
                [](ReductionCandidateGroup &C) { return C.getGEPCount() < 2; }),
            Candidates.end());
    }

    // Sort candidates in descending order from best to worse:
    //  1. First by number of reduced GEP instructions (higher is better).
    //  2. Then by estimated register pressure increase (lower is better).
    auto Comparator = [](ReductionCandidateGroup &L, ReductionCandidateGroup &R)
    {
        if (L.getGEPCount() == R.getGEPCount())
            return L.getScore().RegisterPressure < R.getScore().RegisterPressure;

        return R.getGEPCount() < L.getGEPCount();
    };
    std::sort(Candidates.begin(), Candidates.end(), Comparator);

    // After sorting candidates, iterate over them and do reduction, keeping
    // register pressure under threshold.
    for (auto It = Candidates.begin(); It != Candidates.end(); ++It)
    {
        if (IGC_IS_FLAG_ENABLED(EnableGEPLSRToPreheader) && It->getScore().ReducesInstructions)
        {
            // It is beneficial to reduce to preheader, but keep register pressure in check.
            if (!RPT.fitsPressureThreshold(*It))
            {
                // Above threshold, just simplify index calculation inside loop.
                It->setReductionType(REDUCE_INDEX_ONLY);
            }
        }
        else
        {
            // If nothing from the loop would be reduced to preheader, only
            // simplify index calculation inside loop.
            It->setReductionType(REDUCE_INDEX_ONLY);
        }

        // Check if nothing would be reduced/simplified.
        if (It->getReductionType() == REDUCE_INDEX_ONLY && It->getGEPCount() == 1)
            continue;

        LLVM_DEBUG(
            dbgs() << "  Executing reduction=" << (It->getReductionType() == REDUCE_TO_PREHEADER ? "TO_PREHEADER" : "INDEX_ONLY") << " for ";
            It->print(dbgs());
            dbgs() << "\n");

        It->transform(IRB, E);

        cleanup(*It);
        RPT.updatePressure(*It, E);

        changed = true;
    }

    return changed;
}


void Reducer::cleanup(ReductionCandidateGroup &C)
{
    // Delete GEP instructions together with index calculations. Inform Register
    // Pressure Estimator about removed instructions.
    for (auto *GEP : C.getReduced())
    {
#if LLVM_VERSION_MAJOR < 14
        RecursivelyDeleteTriviallyDeadInstructions(GEP);
#else
        RecursivelyDeleteTriviallyDeadInstructions(GEP, nullptr, nullptr, [&](Value* V) { RPT.trackDeletedInstruction(V); });
#endif
    }
}


// If SCEV is zext/sext, drop extend.
const SCEV *SCEVHelper::dropExt(const SCEV *S)
{
    do
    {
        if (auto *Zext = dyn_cast<SCEVZeroExtendExpr>(S))
            S = Zext->getOperand();
        else if (auto *Sext = dyn_cast<SCEVSignExtendExpr>(S))
            S = Sext->getOperand();
        else
            break;
    } while (true);

    return S;
}


// Returns true is SCEV expression legal.
bool SCEVHelper::isValid(const SCEV *S)
{
    if (isa<SCEVCouldNotCompute>(S))
        return false;

    // Scalar Evolution doesn't have SCEV expression for bitwise-and. Instead,
    // if possible, SE produces expressions for any integer size, leaving cleanup
    // to legalization pass. For example this code:
    //     %1 = shl i64 %0, 32
    //     %2 = ashr exact i64 %1, 30
    // produces i34 integer SCEV.
    //
    // By default don't allow illegal integer types.
    if (IGC_IS_FLAG_ENABLED(EnableGEPLSRAnyIntBitWidth))
        return true;

    std::function<bool(Type*)> IsInvalidInt = [](Type *Ty)
    {
        if (!Ty->isIntegerTy())
            return false;

        auto bits = Ty->getScalarSizeInBits();
        switch (bits)
        {
        case 8:
        case 16:
        case 32:
        case 64:
            return false;
        default:
            return bits > 8;
        }
    };

    bool HasInvalidInt = SCEVExprContains(S, [&](const SCEV *S) {
        if (auto *Cast = dyn_cast<SCEVCastExpr>(S))
            return IsInvalidInt(Cast->getOperand()->getType()) || IsInvalidInt(Cast->getType());
        return false;
    });

    LLVM_DEBUG(
        if (HasInvalidInt) {
            dbgs() << "  Dropping SCEV with invalid integer type: "; S->print(dbgs()); dbgs() << "\n";
        });

    return !HasInvalidInt;
}


bool SCEVHelper::isEqual(const SCEV *A, const SCEV *B)
{
    // Scalar Evolution keeps unique SCEV instances, so we can compare pointers.
    if (A == B)
        return true;

    if (A->getSCEVType() != B->getSCEVType())
        return false;

    switch (A->getSCEVType())
    {
    case scConstant:
        // Can be different bit width, but same integer value.
        return cast<SCEVConstant>(A)->getValue()->getZExtValue() == cast<SCEVConstant>(B)->getValue()->getZExtValue();
    default:
        return false;
    }
}


SCEVHelper::SCEVAddBuilder &SCEVHelper::SCEVAddBuilder::add(const SCEV *S, bool Negative)
{
    IGC_ASSERT(S->getType()->isIntegerTy());

    // strip extend
    S = SCEVHelper::dropExt(S);

    if (auto *Expr = dyn_cast<SCEVAddExpr>(S))
    {
        for (auto *Op : Expr->operands())
            add(Op, Negative);
        return *this;
    }

    Ops.emplace_back(S, Negative);

    return *this;
}


const SCEV *SCEVHelper::SCEVAddBuilder::build()
{
    // ScalarEvolution::getAddExpr requires all operands to have the same
    // type. First find the widest type.
    Type *T = nullptr;
    for (auto *It = Ops.begin(); It != Ops.end(); ++It)
    {
        T = T ? SE.getWiderType(T, It->S->getType()) : It->S->getType();
    }

    // Join list of operands, extending type if required.
    SmallVector<const SCEV*, 16> FinalOps;

    for (auto *It = Ops.begin(); It != Ops.end(); ++It)
    {
        const SCEV *S = It->S;
        S = S->getType() == T ? S : SE.getSignExtendExpr(S, T);
        FinalOps.push_back(It->Negative ? SE.getNegativeSCEV(S) : S);
    }

    return SE.getAddExpr(FinalOps);
}


SCEVHelper::SCEVMulBuilder &SCEVHelper::SCEVMulBuilder::add(const SCEV *S)
{
    IGC_ASSERT(S->getType()->isIntegerTy());

    // strip extend
    S = SCEVHelper::dropExt(S);

    Ops.emplace_back(S);

    return *this;
}


const SCEV *SCEVHelper::SCEVMulBuilder::build()
{
    // ScalarEvolution::getMulExpr requires all operands to have the same
    // type. First find the widest type.
    Type *T = nullptr;
    for (auto S : Ops)
    {
        T = T ? SE.getWiderType(T, S->getType()) : S->getType();
    }

    // Join list of operands, extending type if required.
    SmallVector<const SCEV*, 4> FinalOps;

    for (auto S : Ops)
    {
        FinalOps.push_back(S->getType() == T ? S : SE.getSignExtendExpr(S, T));
    }

    return SE.getMulExpr(FinalOps);
}


GEPLoopStrengthReduction::GEPLoopStrengthReduction(bool AllowLICM)
    : FunctionPass(ID), AllowLICM(AllowLICM)
{
    initializeGEPLoopStrengthReductionPass(*PassRegistry::getPassRegistry());
}


void GEPLoopStrengthReduction::getAnalysisUsage(llvm::AnalysisUsage &AU) const
{
    AU.addRequired<CodeGenContextWrapper>();
    AU.addRequired<DominatorTreeWrapperPass>();
    AU.addRequired<IGCFunctionExternalRegPressureAnalysis>();
    AU.addRequired<IGCLivenessAnalysis>();
    AU.addRequired<llvm::LoopInfoWrapperPass>();
    AU.addRequired<MetaDataUtilsWrapper>();
    AU.addRequired<llvm::ScalarEvolutionWrapperPass>();
}


bool GEPLoopStrengthReduction::runOnFunction(llvm::Function &F)
{
    if (F.hasOptNone())
        return false;

    LLVM_DEBUG(dbgs()
        << "  Running on function " << F.getName()
        << " options: MaxPressureRatio=" << IGC_GET_FLAG_VALUE(GEPLSRThresholdRatio)
        << " EnableLsrToPreheader=" << IGC_GET_FLAG_VALUE(EnableGEPLSRToPreheader) << "\n");

    auto &CGC = *getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
    auto &DL = F.getParent()->getDataLayout();
    auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
    auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
    auto &MMD = *getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
    auto &RPE = getAnalysis<IGCLivenessAnalysis>();
    auto &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();

    // Note: FRPE is a Module analysis and currently runs only once.
    // If function A calls function B then it's possible that transformation of function A
    // increases pressure so LSR reduction should not be applied in function B, but we don't
    // recompute FRPE to save compile time, so reduction might apply for loops in function B.
    auto &FRPE = getAnalysis<IGCFunctionExternalRegPressureAnalysis>();
    auto *WI = &FRPE.getWIAnalysis(&F);

    // Using one SCEV expander between all reductions reduces number of duplicated new instructions.
    auto E = SCEVExpander(SE, DL, "gep-loop-strength-reduction");

    SmallVector<ReductionCandidateGroup, 32> Candidates;

    RegisterPressureTracker RPT(F, CGC, RPE, FRPE, *WI);

    for (Loop *L : LI.getLoopsInPreorder())
        Analyzer(DL, DT, *L, LI, SE, E).analyze(Candidates);

    if (Candidates.empty())
        return false;

    Scorer(DL, MMD, RPE, *WI).score(Candidates);

    IGCLLVM::IRBuilder<> IRB(F.getContext());

    bool changed = Reducer(IRB, RPT, E, AllowLICM).reduce(Candidates);
    if (changed)
        LLVM_DEBUG(dbgs() << "  Modified function " << F.getName() << "\n");

    return changed;
}


// Register pass to igc-opt
#define PASS_FLAG "igc-gep-loop-strength-reduction"
#define PASS_DESCRIPTION "Reduces strength of GEP instructions in loops"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(GEPLoopStrengthReduction, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
IGC_INITIALIZE_PASS_DEPENDENCY(IGCFunctionExternalRegPressureAnalysis)
IGC_INITIALIZE_PASS_DEPENDENCY(IGCLivenessAnalysis)
IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
IGC_INITIALIZE_PASS_END(GEPLoopStrengthReduction, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char GEPLoopStrengthReduction::ID = 0;

FunctionPass *IGC::createGEPLoopStrengthReductionPass(bool AllowLICM)
{
    return new GEPLoopStrengthReduction(AllowLICM);
}
