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

Copyright (C) 2022 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "Compiler/Optimizer/RuntimeValueVectorExtractPass.h"
#include "Compiler/IGCPassSupport.h"
#include "GenISAIntrinsics/GenIntrinsicInst.h"

#include "common/LLVMWarningsPush.hpp"
#include <llvmWrapper/IR/DerivedTypes.h>
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;
using namespace IGC;

#define PASS_FLAG "igc-runtimevalue-vector-extract-pass"
#define PASS_DESCRIPTION "Shader extract element from vector of constants optimization"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(RuntimeValueVectorExtractPass, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_END(RuntimeValueVectorExtractPass, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

namespace IGC {

char RuntimeValueVectorExtractPass::ID = 0;

////////////////////////////////////////////////////////////////////////////
RuntimeValueVectorExtractPass::RuntimeValueVectorExtractPass() : llvm::FunctionPass(ID), changed(false) {
  initializeRuntimeValueVectorExtractPassPass(*llvm::PassRegistry::getPassRegistry());
}

////////////////////////////////////////////////////////////////////////////
void RuntimeValueVectorExtractPass::getAnalysisUsage(llvm::AnalysisUsage &AU) const { AU.setPreservesCFG(); }

////////////////////////////////////////////////////////////////////////////
bool RuntimeValueVectorExtractPass::runOnFunction(llvm::Function &F) {
  changed = false;
  visit(F);
  return changed;
}

////////////////////////////////////////////////////////////////////////////
// @brief Converts extracts of elements from corresponding RuntimeValue
// vector to RuntimeValue calls returning concrete scalars.
// Only extracts using constant indexes are converted.
// Only 32-bit and 64-bit RuntimeValues are supported at the moment.
//
// Replace:
//   %0 = call <8 x i32> @llvm.genx.GenISA.RuntimeValue.v8i32(i32 4)
//   %scalar  = extractelement <8 x i32> %0, i32 0
//   %scalar1 = extractelement <8 x i32> %0, i32 1
// with:
//   %scalar  = call i32 @llvm.genx.GenISA.RuntimeValue.i32(i32 4)
//   %scalar1 = call i32 @llvm.genx.GenISA.RuntimeValue.i32(i32 5)
void RuntimeValueVectorExtractPass::visitExtractElementInst(llvm::ExtractElementInst &I) {
  // Optimization works only on constant indexes
  if (isa<ConstantInt>(I.getIndexOperand())) {
    GenIntrinsicInst *GII = dyn_cast<GenIntrinsicInst>(I.getVectorOperand());
    if (GII && GII->getIntrinsicID() == GenISAIntrinsic::GenISA_RuntimeValue && isa<ConstantInt>(GII->getOperand(0)) &&
        isa<IGCLLVM::FixedVectorType>(GII->getType())) {
      IGCLLVM::FixedVectorType *giiVectorType = cast<IGCLLVM::FixedVectorType>(GII->getType());
      // Only 32-bit and 64-bit values are supported at the moment
      if (giiVectorType->getElementType()->getPrimitiveSizeInBits() == 32 ||
          giiVectorType->getElementType()->getPrimitiveSizeInBits() == 64) {
        bool is64bit = giiVectorType->getElementType()->getPrimitiveSizeInBits() == 64;

        IRBuilder<> Builder(&I);
        Function *runtimeValueFunc = GenISAIntrinsic::getDeclaration(
            I.getModule(), GenISAIntrinsic::GenISA_RuntimeValue, giiVectorType->getElementType());

        const uint32_t eeiIndex = int_cast<uint32_t>(cast<ConstantInt>(I.getIndexOperand())->getZExtValue());
        const uint32_t giiOffset = int_cast<uint32_t>(cast<ConstantInt>(GII->getOperand(0))->getZExtValue());

        // Calculate new offset
        const uint32_t offset = giiOffset + (is64bit ? eeiIndex * 2 : eeiIndex);

        Value *CI = Builder.CreateCall(runtimeValueFunc, Builder.getInt32(offset));

        I.replaceAllUsesWith(CI);
        I.eraseFromParent();

        changed = true;
      } else {
        IGC_ASSERT_MESSAGE(0, "Only 32-bit and 64-bit values are supported at the moment");
      }
    }
  }
}

} // namespace IGC
