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

Copyright (C) 2024 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "CodeGenPublicEnums.h"
#include "Compiler/CodeGenContextWrapper.hpp"
#include "Compiler/CodeGenPublic.h"
#include "Compiler/IGCPassSupport.h"
#include "Compiler/CISACodeGen/helper.h"
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Function.h>
#include "common/LLVMWarningsPop.hpp"
#include "MatchCommonKernelPatterns.hpp"
#include "IGCIRBuilder.h"

using namespace llvm;
using namespace IGC;

char MatchCommonKernelPatterns::ID = 0;

#define INTERPRETER_CASES_NUMBER 18
#define PASS_FLAG "match-common-kernel-patterns"
#define PASS_DESCRIPTION "Match common kernel patterns"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(MatchCommonKernelPatterns, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_END(MatchCommonKernelPatterns, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

MatchCommonKernelPatterns::MatchCommonKernelPatterns() : FunctionPass(ID) {
    initializeMatchCommonKernelPatternsPass(*PassRegistry::getPassRegistry());
}

bool MatchCommonKernelPatterns::runOnFunction(Function &F) {
    Ctx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();

    // Check if F is part of interpreter pattern.
    if (IGC_IS_FLAG_ENABLED(EnableInterpreterPatternMatching))
        if (isInterpreterPattern(F))
            Ctx->m_kernelsWithForcedRetry.push_back(&F);

    return false;
}

bool MatchCommonKernelPatterns::isInterpreterPattern(const Function &F) const {
    // Skip non-kernel functions.
    if (F.getCallingConv() != CallingConv::SPIR_KERNEL)
        return false;

    // Skip kernels with less than 100 basic blocks.
    if (F.size() < 100)
        return false;

    // Iterate over function and check if it has a pattern of Interpreter kernel.
    for (const auto &BB : F)
        if (isBBPartOfInterpreterPattern(&BB))
            return true;

    return false;
}

bool MatchCommonKernelPatterns::isBBPartOfInterpreterPattern(const BasicBlock *BB) const {
    // Try to find BB that is a case of the switch. Try to get a switch argument.
    if (BB->size() != 2)
        return false;

    // Try to find the first switch case.
    Value *CurrentSwitchArg = getInterpreterSwitchArg(BB);
    if (!CurrentSwitchArg)
        return false;

    // Try to find the entrance of the switch.
    const BasicBlock *SwitchEntrance = BB->getSinglePredecessor();
    if (!SwitchEntrance)
        return false;

    // The switch entrance should have conditional branch and the switch argument should be loaded from memory.
    Value *SwitchArg = getInterpreterSwitchArg(SwitchEntrance);
    if (SwitchArg != CurrentSwitchArg)
        return false;

    if (auto Extr = dyn_cast<ExtractElementInst>(SwitchArg)) {
        if (!isa<LoadInst>(Extr->getOperand(0)))
            return false;
    } else if (!isa<LoadInst>(SwitchArg)) {
        return false;
    }

    // Try to find other switch cases.
    size_t Cases = 0;
    for (auto User = SwitchArg->user_begin(), End = SwitchArg->user_end(); User != End; ++User) {
        BasicBlock *CurrentBB = cast<Instruction>(*User)->getParent();
        // Skip the switch entrance and BB that were already checked.
        if (BB == CurrentBB || SwitchEntrance == CurrentBB)
            continue;

        if (getInterpreterSwitchArg(CurrentBB) != SwitchArg)
            continue;

        // Check if CurrentBB has at least one branch which leads to the switch entrance.
        bool HasBackEdge = false;
        for (auto SI = succ_begin(CurrentBB), E = succ_end(CurrentBB); SI != E; ++SI)
            HasBackEdge |= isBBBackEdge(*SI, SwitchEntrance);

        if (HasBackEdge)
            Cases++;

        // If switch has more than INTERPRETER_CASES_NUMBER cases then it is a Interpreter pattern.
        if (Cases == INTERPRETER_CASES_NUMBER)
            return true;
    }

    return false;
}

bool MatchCommonKernelPatterns::isBBBackEdge(const BasicBlock *BB, const BasicBlock *EntryBB) const {
    // Try to find critical edge or back edge.
    const BasicBlock *CritOrBackEdge = BB->getSingleSuccessor();
    if (!CritOrBackEdge)
        return false;

    if (CritOrBackEdge == EntryBB)
        return true;

    return CritOrBackEdge->getSingleSuccessor() == EntryBB;
}

Value *MatchCommonKernelPatterns::getInterpreterSwitchArg(const BasicBlock *BB) const {
    const BranchInst *Br = dyn_cast<const BranchInst>(BB->getTerminator());
    if (!Br || !Br->isConditional())
        return nullptr;

    const Value *Condition = Br->getCondition();
    if (!isa<ICmpInst>(Condition))
        return nullptr;

    const ICmpInst *Cmp = cast<const ICmpInst>(Condition);
    Value *NonConstOperand = nullptr;
    if (isa<ConstantInt>(Cmp->getOperand(0))) {
        NonConstOperand = Cmp->getOperand(1);
    } else if (isa<ConstantInt>(Cmp->getOperand(1))) {
        NonConstOperand = Cmp->getOperand(0);
    } else {
        return nullptr;
    }

    return NonConstOperand;
}