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

Copyright (C) 2023-2026 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "HandleSpirvDecorationMetadata.h"
#include "AdaptorOCL/Utils/CacheControlsHelper.h"
#include "Compiler/IGCPassSupport.h"
#include "Compiler/CISACodeGen/OpenCLKernelCodeGen.hpp"
#include "GenISAIntrinsics/GenIntrinsicInst.h"

#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/Module.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/Demangle/Demangle.h>
#include <llvm/IR/Mangler.h>
#include <llvm/Support/Regex.h>
#include "common/LLVMWarningsPop.hpp"
#include "llvmWrapper/IR/DerivedTypes.h"
#include "llvmWrapper/IR/Instructions.h"

#include "spirv/unified1/spirv.hpp"

#include <regex>

using namespace llvm;
using namespace IGC;

// Register pass to igc-opt
#define PASS_FLAG "igc-handle-spirv-decoration-metadata"
#define PASS_DESCRIPTION "Handle spirv.Decoration metadata"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(HandleSpirvDecorationMetadata, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_END(HandleSpirvDecorationMetadata, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char HandleSpirvDecorationMetadata::ID = 0;

HandleSpirvDecorationMetadata::HandleSpirvDecorationMetadata() : ModulePass(ID) {
  initializeHandleSpirvDecorationMetadataPass(*PassRegistry::getPassRegistry());
}

// Get type from demangled name for both typed and opaque pointers
// For opaque pointers, we cannot deduce type from ptr
Type *HandleSpirvDecorationMetadata::getArgumentType(std::string demangledName, size_t argNo, llvm::LLVMContext &ctx) {
  auto open = demangledName.find('('), close = demangledName.find(')', open);
  if (open == std::string::npos || close == std::string::npos || close <= open + 1) {
    IGC_ASSERT_MESSAGE(0, "Unexpected demangled name!");
    return nullptr;
  }

  std::stringstream ss(demangledName.substr(open + 1, close - open - 1));
  std::string arg;
  std::vector<std::string> args;

  while (std::getline(ss, arg, ',')) {
    arg.erase(0, arg.find_first_not_of(" \t"));
    arg.erase(arg.find_last_not_of(" \t") + 1);
    args.push_back(arg);
  }

  if (argNo >= args.size()) {
    IGC_ASSERT_MESSAGE(0, "Provided arg number is greater than or equal to arg.size()!");
    return nullptr;
  }

  const auto &a = args[argNo];

  if (std::regex_search(a, std::regex("\\bchar\\b")))
    return llvm::Type::getInt8Ty(ctx);
  else if (std::regex_search(a, std::regex("\\bshort\\b")))
    return llvm::Type::getInt16Ty(ctx);
  else if (std::regex_search(a, std::regex("\\bint\\b")))
    return llvm::Type::getInt32Ty(ctx);
  else if (std::regex_search(a, std::regex("\\blong\\b")))
    return llvm::Type::getInt64Ty(ctx);
  else
    IGC_ASSERT_MESSAGE(0, "Unexpected type!");

  return nullptr;
}

bool HandleSpirvDecorationMetadata::runOnModule(Module &module) {
  m_Metadata = getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
  m_pCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
  m_Module = &module;
  m_BuiltinsToRemove.clear();

  handleInstructionsDecorations();
  handleGlobalVariablesDecorations();

  for (auto &F : m_BuiltinsToRemove)
    F->eraseFromParent();

  return m_changed;
}

void HandleSpirvDecorationMetadata::handleInstructionsDecorations() { visit(m_Module); }

void HandleSpirvDecorationMetadata::handleGlobalVariablesDecorations() {
  for (auto &globalVariable : m_Module->globals()) {
    auto spirvDecorations = parseSPIRVDecorationsFromMD(&globalVariable);
    for (auto &[DecorationId, MDNodes] : spirvDecorations) {
      switch (DecorationId) {
      // IDecHostAccessINTEL
      // We use two IDs (6147 and 6188) to preserve backward compatibility with
      // older versions
      case 6147: // Old ID
      case 6188: // New ID
      {
        IGC_ASSERT_MESSAGE(MDNodes.size() == 1, "Only one HostAccessINTEL decoration can be applied "
                                                "to a single global variable!");
        handleHostAccessIntel(globalVariable, *MDNodes.begin());
        break;
      }
      default:
        continue;
      }
    }
  }
}

void HandleSpirvDecorationMetadata::handleHostAccessIntel(GlobalVariable &globalVariable, MDNode *node) {
  globalVariable.addAttribute("host_var_name", dyn_cast<MDString>(node->getOperand(2))->getString());

  m_changed = true;
  m_Metadata->capabilities.globalVariableDecorationsINTEL = true;
}

void HandleSpirvDecorationMetadata::visitBinaryOperator(BinaryOperator &I) {
  if (I.getOpcode() != Instruction::FDiv)
    return;

  auto SpirvDecorations = parseSPIRVDecorationsFromMD(&I);
  for (auto &[DecorationId, MDNodes] : SpirvDecorations) {
    switch (DecorationId) {
    case spv::DecorationFPRoundingMode: {
      handleFPRoundingMode(I, MDNodes);
      // FDiV is replaced and should no longer be operated on
      return;
    }
    default:
      continue;
    }
  }
}

void HandleSpirvDecorationMetadata::visitLoadInst(LoadInst &I) {
  auto spirvDecorations = parseSPIRVDecorationsFromMD(I.getPointerOperand());
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTEL<LoadCacheControl>(I, MDNodes);
      break;
    }
    default:
      continue;
    }
  }
}

void HandleSpirvDecorationMetadata::visitStoreInst(StoreInst &I) {
  auto spirvDecorations = parseSPIRVDecorationsFromMD(I.getPointerOperand());
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlStoreINTEL
    case DecorationIdCacheControlStore: {
      handleCacheControlINTEL<StoreCacheControl>(I, MDNodes);
      break;
    }
    default:
      continue;
    }
  }
}

void HandleSpirvDecorationMetadata::visit2DBlockReadCallInst(CallInst &I, StringRef unmangledName) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTELFor2DBlockIO<LoadCacheControl>(I, MDNodes, unmangledName);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visit2DBlockWriteCallInst(CallInst &I, StringRef unmangledName) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlStoreINTEL
    case DecorationIdCacheControlStore: {
      handleCacheControlINTELFor2DBlockIO<StoreCacheControl>(I, MDNodes, unmangledName);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visitPrefetchCallInst(CallInst &I) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTELForPrefetch(I, MDNodes);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visit1DBlockReadCallInst(CallInst &I) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTELFor1DBlockIO<LoadCacheControl>(I, MDNodes);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visit1DBlockWriteCallInst(CallInst &I) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlStoreINTEL
    case DecorationIdCacheControlStore: {
      handleCacheControlINTELFor1DBlockIO<StoreCacheControl>(I, MDNodes);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visit1DBlockPrefetchCallInst(CallInst &I) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTELFor1DBlockIO<LoadCacheControl>(I, MDNodes);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visitOCL1DBlockPrefetchCallInst(CallInst &I, SmallVectorImpl<StringRef> &Matches) {
  Value *ptr = I.getArgOperand(0);
  auto spirvDecorations = parseSPIRVDecorationsFromMD(ptr);
  for (auto &[DecorationId, MDNodes] : spirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTELForOCL1DBlockPrefetch(I, MDNodes, Matches);
      break;
    }
    }
  }
}

void HandleSpirvDecorationMetadata::visitPredicatedLoadInst(CallInst &I) {
  auto SpirvDecorations = parseSPIRVDecorationsFromMD(I.getArgOperand(0));
  for (auto &[DecorationId, MDNodes] : SpirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlLoadINTEL
    case DecorationIdCacheControlLoad: {
      handleCacheControlINTEL<LoadCacheControl>(I, MDNodes);
      break;
    }
    default:
      continue;
    }
  }
}

void HandleSpirvDecorationMetadata::visitPredicatedStoreInst(CallInst &I) {
  auto SpirvDecorations = parseSPIRVDecorationsFromMD(I.getArgOperand(0));
  for (auto &[DecorationId, MDNodes] : SpirvDecorations) {
    switch (DecorationId) {
    // IDecCacheControlStoreINTEL
    case DecorationIdCacheControlStore: {
      handleCacheControlINTEL<StoreCacheControl>(I, MDNodes);
      break;
    }
    default:
      continue;
    }
  }
}

void HandleSpirvDecorationMetadata::visitCallInst(CallInst &I) {
  Function *F = I.getCalledFunction();
  if (!F)
    return;

  static const Regex pattern2DBlockRead(
      "_Z[0-9]+(intel_sub_group_2d_block_(prefetch|read|read_transform|read_transpose)_[0-9]+b_[0-9]+r[0-9]+x[0-9]+c)");
  static const Regex pattern2DBlockWrite("_Z[0-9]+(intel_sub_group_2d_block_write_[0-9]+b_[0-9]+r[0-9]+x[0-9]+c)");
  static const Regex patternOCL1DBlockPrefetch("_Z[0-9]+(intel_sub_group_block_prefetch_(uc|us|ui|ul)(2|4|8|16)?)");
  static const Regex patternPrefetch("_Z[0-9]+__spirv_ocl_prefetch");
  static const Regex pattern1DBlockRead("_Z[0-9]+__spirv_SubgroupBlockReadINTEL");
  static const Regex pattern1DBlockWrite("_Z[0-9]+__spirv_SubgroupBlockWriteINTEL");
  static const Regex pattern1DBlockPrefetch("_Z[0-9]+__spirv_SubgroupBlockPrefetchINTEL");
  static const Regex patternPredicatedReadSPV("_Z[0-9]+__spirv_PredicatedLoadINTEL.+");
  static const Regex patternPredicatedWriteSPV("_Z[0-9]+__spirv_PredicatedStoreINTEL.+");
  static const Regex patternSpirvSqrt("_Z[0-9]+__spirv_ocl_sqrt(f|d|Dh)");

  SmallVector<StringRef, 4> Matches;
  StringRef FuncName = F->getName();

  if (pattern2DBlockRead.match(FuncName, &Matches)) {
    visit2DBlockReadCallInst(I, Matches[1]);
  } else if (pattern2DBlockWrite.match(FuncName, &Matches)) {
    visit2DBlockWriteCallInst(I, Matches[1]);
  } else if (patternPrefetch.match(FuncName, &Matches)) {
    visitPrefetchCallInst(I);
  } else if (pattern1DBlockRead.match(FuncName, &Matches)) {
    visit1DBlockReadCallInst(I);
  } else if (pattern1DBlockWrite.match(FuncName, &Matches)) {
    visit1DBlockWriteCallInst(I);
  } else if (pattern1DBlockPrefetch.match(FuncName, &Matches)) {
    visit1DBlockPrefetchCallInst(I);
  } else if (patternOCL1DBlockPrefetch.match(FuncName, &Matches)) {
    visitOCL1DBlockPrefetchCallInst(I, Matches);
  } else if (patternPredicatedReadSPV.match(FuncName)) {
    visitPredicatedLoadInst(I);
  } else if (patternPredicatedWriteSPV.match(FuncName)) {
    visitPredicatedStoreInst(I);
  } else if (patternSpirvSqrt.match(FuncName)) {
    visitSpirvSqrtCallInst(I);
  }
}

template <typename T>
void HandleSpirvDecorationMetadata::handleCacheControlINTEL(Instruction &I, SmallPtrSetImpl<MDNode *> &MDNodes) {
  static_assert(std::is_same_v<T, LoadCacheControl> || std::is_same_v<T, StoreCacheControl>);
  CacheControlFromMDNodes cacheControl = resolveCacheControlFromMDNodes<T>(m_pCtx, MDNodes);
  if (cacheControl.isEmpty)
    return;
  if (cacheControl.isInvalid) {
    m_pCtx->EmitWarning("Unsupported cache controls configuration requested. Applying default configuration.", &I);
    return;
  }

  MDNode *CacheCtrlNode = MDNode::get(
      I.getContext(), ConstantAsMetadata::get(ConstantInt::get(Type::getInt32Ty(I.getContext()), cacheControl.value)));
  I.setMetadata("lsc.cache.ctrl", CacheCtrlNode);
  m_changed = true;
}

template <typename T>
void HandleSpirvDecorationMetadata::handleCacheControlINTELFor2DBlockIO(CallInst &I, SmallPtrSetImpl<MDNode *> &MDNodes,
                                                                        StringRef unmangledName) {
  static_assert(std::is_same_v<T, LoadCacheControl> || std::is_same_v<T, StoreCacheControl>);
  CacheControlFromMDNodes cacheControl = resolveCacheControlFromMDNodes<T>(m_pCtx, MDNodes);
  if (cacheControl.isEmpty)
    return;
  if (cacheControl.isInvalid) {
    m_pCtx->EmitWarning("Unsupported cache controls configuration requested. Applying default configuration.", &I);
    return;
  }

  Function *F = I.getCalledFunction();
  IGC_ASSERT(F);

  SmallVector<Value *, 4> args(I.args());
  args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()), cacheControl.value));

  SmallVector<Type *, 4> argTypes;
  for (const auto &arg : args)
    argTypes.push_back(arg->getType());

  FunctionType *FT = FunctionType::get(I.getType(), argTypes, false);
  std::string newFuncName = "__internal_" + unmangledName.str() + "_cache_controls";

  auto newFunction = m_Module->getOrInsertFunction(newFuncName, FT);
  auto newCall = IGCLLVM::createCallInst(newFunction, args, "", &I);

  I.replaceAllUsesWith(newCall);
  I.eraseFromParent();
  m_changed = true;

  // Cleanup unused function if all calls have been replaced with the internal version
  if (F->getNumUses() == 0)
    m_BuiltinsToRemove.insert(F);
}

void HandleSpirvDecorationMetadata::handleCacheControlINTELForPrefetch(llvm::CallInst &I,
                                                                       llvm::SmallPtrSetImpl<llvm::MDNode *> &MDNodes) {
  CacheControlFromMDNodes cacheControl = resolveCacheControlFromMDNodes<LoadCacheControl>(m_pCtx, MDNodes);
  if (cacheControl.isEmpty)
    return;
  if (cacheControl.isInvalid) {
    m_pCtx->EmitWarning("Unsupported cache controls configuration requested. Applying default configuration.", &I);
    return;
  }

  Function *F = I.getCalledFunction();
  IGC_ASSERT(F);

  // Convert prefetch call to: __lsc_prefetch_cache_controls(global void* p, int element_size, int num_elements, enum
  // LSC_LDCC cache_opt)
  SmallVector<Value *, 4> args;
  args.push_back(I.getArgOperand(0));

  // OpenCL spec states for prefetch: "Prefetch num_gentypes * sizeof(gentype) bytes into the global cache.".
  // This design is not friendly to opaque pointers, as it assumes element size can be read from pointer.
  // For now read size from typed pointer, and in future this will be replaced with opaque prefetch with
  // explicit element size as arg.
  PointerType *PTy = dyn_cast<PointerType>(I.getArgOperand(0)->getType());
  IGC_ASSERT(PTy);

  // As of today there's no "OpUntypedPrefetch" which as previous comments mentions,
  // is needed to implement this prefetch correctly on opaque pointers.
  // It will of course result in performance penalty and needs to be changed once
  // "OpUntypedPrefetch" is ready.
  if (IGCLLVM::isPointerTy(PTy))
    return;

  args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()),
                                  IGCLLVM::getNonOpaquePtrEltTy(PTy)->getPrimitiveSizeInBits() / 8));

  // OpenCL prefetch overloads num_elements to either i32 or i64. Convert to i32.
  IGCLLVM::IRBuilder<> builder(&I);
  args.push_back(builder.CreateZExtOrTrunc(I.getArgOperand(1), Type::getInt32Ty(I.getContext())));

  auto config = supportedLoadConfigs.find(static_cast<LSC_L1_L3_CC>(cacheControl.value));
  if (m_pCtx->platform.getPlatformInfo().eProductFamily == IGFX_PVC && config != supportedLoadConfigs.end() &&
      config->second.L1 == LoadCacheControl::Cached) {
    m_pCtx->EmitWarning("Prefetch to L1 is unsupported on this platform.", &I);
    args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()),
                                    mapToLSCCacheControl(LoadCacheControl::Uncached, config->second.L3)));
  } else {
    args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()), cacheControl.value));
  }

  SmallVector<Type *, 4> argTypes;
  for (const auto &arg : args)
    argTypes.push_back(arg->getType());

  FunctionType *FT = FunctionType::get(I.getType(), argTypes, false);
  std::string newFuncName = "__lsc_prefetch_cache_controls";

  auto newFunction = m_Module->getOrInsertFunction(newFuncName, FT);
  auto newCall = IGCLLVM::createCallInst(newFunction, args, "", &I);

  I.replaceAllUsesWith(newCall);
  I.eraseFromParent();
  m_changed = true;

  // Cleanup unused function if all calls have been replaced with the internal version
  if (F->getNumUses() == 0)
    m_BuiltinsToRemove.insert(F);
}

template <typename T>
void HandleSpirvDecorationMetadata::handleCacheControlINTELFor1DBlockIO(CallInst &I,
                                                                        SmallPtrSetImpl<MDNode *> &MDNodes) {
  static_assert(std::is_same_v<T, LoadCacheControl> || std::is_same_v<T, StoreCacheControl>);
  CacheControlFromMDNodes cacheControl = resolveCacheControlFromMDNodes<T>(m_pCtx, MDNodes);
  if (cacheControl.isEmpty)
    return;
  if (cacheControl.isInvalid) {
    m_pCtx->EmitWarning("Unsupported cache controls configuration requested. Applying default configuration.", &I);
    return;
  }

  Function *F = I.getCalledFunction();
  IGC_ASSERT(F);

  Type *operationType = nullptr;
  std::string funcName;
  if constexpr (std::is_same_v<T, LoadCacheControl>) {
    if (auto isPrefetch = I.getType()->isVoidTy()) {
      operationType = getArgumentType(llvm::demangle(F->getName().str()), 0, F->getContext());
      funcName = "SubgroupBlockPrefetchINTEL";
    } else {
      operationType = I.getType();
      funcName = "SubgroupBlockReadINTEL";
    }
  } else {
    operationType = I.getArgOperand(1)->getType();
    funcName = "SubgroupBlockWriteINTEL";
  }

  std::string typeName;
  uint32_t numElements = 1;
  Type *elementType = operationType;
  if (auto *vecTy = dyn_cast<IGCLLVM::FixedVectorType>(operationType)) {
    numElements = (uint32_t)vecTy->getNumElements();
    elementType = vecTy->getElementType();
  }

  if (elementType->isIntegerTy()) {
    switch (elementType->getIntegerBitWidth()) {
    case 8:
      typeName = "char";
      break;
    case 16:
      typeName = "short";
      break;
    case 32:
      typeName = "int";
      break;
    case 64:
      typeName = "long";
      break;
    default:
      IGC_ASSERT(0 && "Unsupported integer type");
      break;
    }
  }

  if (numElements > 1)
    typeName += std::to_string(numElements);

  SmallVector<Value *, 3> args(I.args());
  args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()), cacheControl.value));

  SmallVector<Type *, 3> argTypes;
  for (const auto &arg : args)
    argTypes.push_back(arg->getType());

  FunctionType *funcTy = FunctionType::get(I.getType(), argTypes, false);
  std::string newFuncName = "__internal_" + funcName + "_" + typeName + "_cache_controls";

  auto newFunction = m_Module->getOrInsertFunction(newFuncName, funcTy);
  auto newCall = IGCLLVM::createCallInst(newFunction, args, "", &I);

  I.replaceAllUsesWith(newCall);
  I.eraseFromParent();
  m_changed = true;

  // Cleanup unused function if all calls have been replaced with the internal version
  if (F->getNumUses() == 0)
    m_BuiltinsToRemove.insert(F);
}

void HandleSpirvDecorationMetadata::handleCacheControlINTELForOCL1DBlockPrefetch(CallInst &I,
                                                                                 SmallPtrSetImpl<MDNode *> &MDNodes,
                                                                                 SmallVectorImpl<StringRef> &Matches) {
  IGC_ASSERT(IGCLLVM::starts_with(Matches[1], "intel_sub_group_block_prefetch"));

  CacheControlFromMDNodes cacheControl = resolveCacheControlFromMDNodes<LoadCacheControl>(m_pCtx, MDNodes);
  if (cacheControl.isEmpty)
    return;
  if (cacheControl.isInvalid) {
    m_pCtx->EmitWarning("Unsupported cache controls configuration requested. Applying default configuration.", &I);
    return;
  }

  Function *F = I.getCalledFunction();
  IGC_ASSERT(F);

  StringRef numElementsFromName = Matches[3] != "" ? Matches[3] : "1";
  uint32_t numElementsToPrefetch = std::stoi(numElementsFromName.str());
  IGC_ASSERT(numElementsToPrefetch == 1 || numElementsToPrefetch == 2 || numElementsToPrefetch == 4 ||
             numElementsToPrefetch == 8 || numElementsToPrefetch == 16);

  uint32_t typeSizeInBytes = 0;
  if (Matches[2] == "uc")
    typeSizeInBytes = 1;
  else if (Matches[2] == "us")
    typeSizeInBytes = 2;
  else if (Matches[2] == "ui")
    typeSizeInBytes = 4;
  else if (Matches[2] == "ul")
    typeSizeInBytes = 8;
  else
    IGC_ASSERT(0 && "Unsupported type prefetch!");

  std::string typeName;
  switch (typeSizeInBytes) {
  case 1:
    typeName = "char";
    break;
  case 2:
    typeName = "short";
    break;
  case 4:
    typeName = "int";
    break;
  case 8:
    typeName = "long";
    break;
  default:
    IGC_ASSERT(0 && "Unsupported block prefetch!");
    break;
  }

  Value *numBytesArg = (ConstantInt::get(Type::getInt32Ty(I.getContext()), (typeSizeInBytes * numElementsToPrefetch)));

  SmallVector<Value *, 3> args(I.args());
  args.push_back(numBytesArg);
  args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()), cacheControl.value));

  SmallVector<Type *, 3> argTypes;
  for (const auto &arg : args)
    argTypes.push_back(arg->getType());

  auto *funcTy = FunctionType::get(I.getType(), argTypes, false);
  auto newFuncName = "__internal_SubgroupBlockPrefetchINTEL_" + typeName + "_cache_controls";

  auto newFunction = m_Module->getOrInsertFunction(newFuncName, funcTy);
  auto newCall = IGCLLVM::createCallInst(newFunction, args, "", &I);

  I.replaceAllUsesWith(newCall);
  I.eraseFromParent();
  m_changed = true;

  // Cleanup unused function if all calls have been replaced with the internal version
  if (F->getNumUses() == 0)
    m_BuiltinsToRemove.insert(F);
}

void HandleSpirvDecorationMetadata::visitSpirvSqrtCallInst(CallInst &I) {
  auto SpirvDecorations = parseSPIRVDecorationsFromMD(&I);
  for (auto &[DecorationId, MDNodes] : SpirvDecorations) {
    switch (DecorationId) {
    case spv::DecorationFPRoundingMode: {
      handleFPRoundingMode(I, MDNodes);
      // call is replaced and should no longer be operated on
      return;
    }
    default:
      continue;
    }
  }
}

void HandleSpirvDecorationMetadata::handleFPRoundingMode(Instruction &I, SmallPtrSetImpl<MDNode *> &MDNodes) {
  IGC_ASSERT_MESSAGE(MDNodes.size() == 1, "Only one FPRoundingMode decoration can be applied to a single instruction!");

  MDNode *Node = *MDNodes.begin();
  // Decoration node layout: operand(0) = decoration ID (39), operand(1) = rounding mode value
  IGC_ASSERT(Node->getNumOperands() >= 2);

  auto *RoundingModeMD = dyn_cast<ValueAsMetadata>(Node->getOperand(1));
  IGC_ASSERT(RoundingModeMD);
  auto *RoundingModeCI = dyn_cast<ConstantInt>(RoundingModeMD->getValue());
  IGC_ASSERT(RoundingModeCI);

  uint64_t RoundingMode = RoundingModeCI->getZExtValue();

  // SPIR-V FPRoundingMode values:
  // 0 = RTE (Round to Nearest Even)
  // 1 = RTZ (Round toward Zero)
  // 2 = RTP (Round toward Positive infinity)
  // 3 = RTN (Round toward Negative infinity)
  IGC_ASSERT_MESSAGE(RoundingMode <= 3, "Invalid FPRoundingMode value!");

  // Convert SPIR-V FPRoundingMode to internal ERoundingMode:
  // internal: 0=ROUND_TO_NEAREST_EVEN, 1=ROUND_TO_POSITIVE, 2=ROUND_TO_NEGATIVE, 3=ROUND_TO_ZERO
  constexpr uint64_t SpirvToInternal[] = {0, 3, 1, 2};
  uint64_t InternalRM = SpirvToInternal[RoundingMode];

  Type *OpType = I.getType();

  // Only fp16, fp32, and fp64 are supported
  if (!OpType->isHalfTy() && !OpType->isFloatTy() && !OpType->isDoubleTy())
    return;

  // For half types, emulate: fpext to float -> compute in RNE -> ftof back to half
  bool IsHalf = OpType->isHalfTy();
  Type *ComputeType = IsHalf ? Type::getFloatTy(I.getContext()) : OpType;
  uint64_t ComputeRM = IsHalf ? 0 /* RNE */ : InternalRM;

  IGCLLVM::IRBuilder<> Builder(&I);

  GenISAIntrinsic::ID IntrinsicId;
  SmallVector<Value *, 3> Args;
  std::string Name;
  if (auto *BinOp = dyn_cast<BinaryOperator>(&I)) {
    // FDiv -> GenISA_IEEE_Divide_rm
    IntrinsicId = GenISAIntrinsic::GenISA_IEEE_Divide_rm;
    Value *Op0 = BinOp->getOperand(0);
    Value *Op1 = BinOp->getOperand(1);
    Args.push_back(IsHalf ? Builder.CreateFPExt(Op0, ComputeType) : Op0);
    Args.push_back(IsHalf ? Builder.CreateFPExt(Op1, ComputeType) : Op1);
    Name = "div_rm";
  } else if (auto *CI = dyn_cast<CallInst>(&I)) {
    // sqrt -> GenISA_IEEE_Sqrt_rm
    IntrinsicId = GenISAIntrinsic::GenISA_IEEE_Sqrt_rm;
    Value *Op0 = CI->getArgOperand(0);
    Args.push_back(IsHalf ? Builder.CreateFPExt(Op0, ComputeType) : Op0);
    Name = "sqrt_rm";
  } else {
    return;
  }

  Args.push_back(ConstantInt::get(Type::getInt32Ty(I.getContext()), ComputeRM));

  Function *IntrinsicFunc = GenISAIntrinsic::getDeclaration(m_Module, IntrinsicId, ComputeType);
  auto *NewCall = IGCLLVM::createCallInst(IntrinsicFunc, Args, Name, &I);
  NewCall->copyFastMathFlags(&I);
  NewCall->setDebugLoc(I.getDebugLoc());

  // For half: truncate float result back to half with the desired rounding mode
  Value *Result = NewCall;
  if (IsHalf) {
    constexpr GenISAIntrinsic::ID SpirvRMToFtof[] = {
        GenISAIntrinsic::GenISA_ftof_rte, // 0 = RTE
        GenISAIntrinsic::GenISA_ftof_rtz, // 1 = RTZ
        GenISAIntrinsic::GenISA_ftof_rtp, // 2 = RTP
        GenISAIntrinsic::GenISA_ftof_rtn  // 3 = RTN
    };
    Type *OverloadTypes[] = {OpType, ComputeType};
    Function *FtofFunc = GenISAIntrinsic::getDeclaration(m_Module, SpirvRMToFtof[RoundingMode], OverloadTypes);
    Result = IGCLLVM::createCallInst(FtofFunc, {NewCall}, "ftof_rm", &I);
    cast<Instruction>(Result)->setDebugLoc(I.getDebugLoc());
  }

  // Save original function before erasing the instruction
  Function *OrigFunc = nullptr;
  if (auto *CI = dyn_cast<CallInst>(&I))
    OrigFunc = CI->getCalledFunction();

  I.replaceAllUsesWith(Result);
  I.eraseFromParent();
  m_changed = true;

  if (OrigFunc && OrigFunc->getNumUses() == 0)
    m_BuiltinsToRemove.insert(OrigFunc);
}
