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

Copyright (C) 2021-2023 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "ModuleAllocaAnalysis.hpp"
#include "Compiler/Optimizer/OpenCLPasses/KernelArgs/KernelArgs.hpp"
#include "Compiler/MetaDataUtilsWrapper.h"
#include "Compiler/IGCPassSupport.h"
#include "Compiler/CISACodeGen/GenCodeGenModule.h"
#include "llvm/IR/Instructions.h"

using namespace llvm;

namespace IGC {
ModuleAllocaAnalysis::ModuleAllocaAnalysis() : ModulePass(ID) {
  initializeModuleAllocaAnalysisPass(*PassRegistry::getPassRegistry());
}

ModuleAllocaAnalysis::~ModuleAllocaAnalysis() {
  for (auto I = InfoMap.begin(), E = InfoMap.end(); I != E; ++I)
    delete I->second;
}

void ModuleAllocaAnalysis::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.setPreservesAll();
  AU.addRequired<MetaDataUtilsWrapper>();
  AU.addRequired<CodeGenContextWrapper>();
}

StringRef ModuleAllocaAnalysis::getPassName() const { return "ModuleAllocaAnalysis"; }

bool ModuleAllocaAnalysis::runOnModule(Module &mod) {
  M = &mod;
  FGA = getAnalysisIfAvailable<GenXFunctionGroupAnalysis>();
  analyze();

  return false;
}

bool ModuleAllocaAnalysis::safeToUseScratchSpace() const {
  IGC_ASSERT(M);

  IGCMD::MetaDataUtils *pMdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
  ModuleMetaData &modMD = *getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
  CodeGenContext &Ctx = *getAnalysis<CodeGenContextWrapper>().getCodeGenContext();

  //
  // Update UseScratchSpacePrivateMemory based on WA and be consistent with
  // the implementation of CEncoder::ByteScatter().
  //
  if (Ctx.m_DriverInfo.NeedWAToTransformA32MessagesToA64() &&
      Ctx.platform.getWATable().WaNoA32ByteScatteredStatelessMessages) {
    return false;
  }

  //
  // For now, all APIs that use scratch space for private memory, must use scratch
  // memory except OpenCL, which can also use non-scratch space. For debugging
  // purpose, a registry key is used for OCL to turn ocl-use-scratch on/off.
  //
  bool supportsScratchSpacePrivateMemory = Ctx.m_DriverInfo.supportsScratchSpacePrivateMemory();
  bool supportsStatelessSpacePrivateMemory = Ctx.m_DriverInfo.supportsStatelessSpacePrivateMemory();
  bool bOCLLegacyStatelessCheck = true;

  if (supportsScratchSpacePrivateMemory) {
    if (Ctx.type == ShaderType::OPENCL_SHADER) {
      if (Ctx.platform.hasScratchSurface() && !Ctx.m_DriverInfo.UseScratchSpaceForATSPlus()) {
        supportsScratchSpacePrivateMemory = Ctx.platform.useScratchSpaceForOCL();
        // IGC has some legacy cases where stateless private memory must be used. This flag is to remove them. If
        // regression happens, revert it.
        bOCLLegacyStatelessCheck =
            !(Ctx.platform.hasScratchSurface() && IGC_IS_FLAG_ENABLED(RemoveLegacyOCLStatelessPrivateMemoryCases));
      }
    }
  }

  if (Ctx.allocatePrivateAsGlobalBuffer()) {
    return false;
  }

  if ((modMD.compOpt.OptDisable && bOCLLegacyStatelessCheck) || !supportsScratchSpacePrivateMemory) {
    return false;
  }

  //
  // Do not use scratch space if module has any stack call.
  // Do not use scratch space if modeule has any variable length alloca
  //
  if (bOCLLegacyStatelessCheck) {
    if (auto *FGA = getAnalysisIfAvailable<GenXFunctionGroupAnalysis>()) {
      if (FGA->getModule() == M) {
        for (auto &I : *FGA) {
          if (FGA->isIndirectCallGroup(I) && !I->isSingle())
            return false;
          if (I->hasStackCall() || I->hasVariableLengthAlloca())
            return false;
        }
      }
    } else {
      // Check individual functions if FGA not available
      for (auto &F : *M) {
        if (F.hasFnAttribute("visaStackCall") || F.hasFnAttribute("hasVLA"))
          return false;
      }
    }
  }

  const llvm::DataLayout *DL = &M->getDataLayout();

  for (auto &F : *M) {
    if (F.isDeclaration())
      continue;

    // Check each instr of this function.
    for (auto &BB : F) {
      for (auto &I : BB) {
        if (AddrSpaceCastInst *CI = dyn_cast<AddrSpaceCastInst>(&I)) {
          // It is not safe to use scratch space as private memory if kernel does
          // AS casting to ADDRESS_SPACE_GLOBAL_OR_PRIVATE or ADDRESS_SPACE_PRIVATE.
          // See speical hack CI code generated at ProgramScopeConstantResolution
          const ADDRESS_SPACE targetAS = (ADDRESS_SPACE)(cast<PointerType>(CI->getType()))->getAddressSpace();
          if (targetAS == ADDRESS_SPACE_GLOBAL_OR_PRIVATE || targetAS == ADDRESS_SPACE_PRIVATE) {
            return false;
          }
        }
        if (Ctx.type == ShaderType::OPENCL_SHADER) {
          // PtrToInt may be used to test if pointer is null, then we cannot
          // distinguish nullptr versus zero offset. This causes a problem
          // with an OpenCL3.0 test. see cassian/oclc_address_space_qualifiers
          if (PtrToIntInst *IPI = dyn_cast<PtrToIntInst>(&I)) {
            if (IPI->getPointerAddressSpace() == ADDRESS_SPACE_PRIVATE) {
              return false;
            }
          }
        }
      }
    }

    if (!isEntryFunc(pMdUtils, &F))
      continue;

    //
    // OCL kernel arguments with type like queue_t and struct are expressed as
    // pointer type. Since there is no explicit AS associated with those pointers,
    // e.g., %opencl.queue_t*, to have both host and device use the same pointer
    // size for those arguments, it is better to disable the use of scratch memory.
    //
    // TODO: fixed those types (they should be in global address space)
    if (Ctx.type == ShaderType::OPENCL_SHADER && IGC_IS_FLAG_ENABLED(ForceStatelessForQueueT)) {
      if (!F.arg_empty()) {
        KernelArgs kernelArgs(F, DL, pMdUtils, &modMD, Ctx.platform.getGRFSize());
        for (const auto &arg : kernelArgs) {
          const KernelArg::ArgType argTy = arg.getArgType();
          if (argTy == KernelArg::ArgType::PTR_DEVICE_QUEUE) {
            return false;
          }
        }
      }
    }

    //
    // Each thread has up to 2 MB scratch space to use. That is, each WI
    // has up to (2*1024*1024 / 8) bytes of scratch space in SIMD8 mode.
    //
    auto funcInfoMD = pMdUtils->getFunctionsInfoItem(&F);
    bool isGeometryStageShader = Ctx.type == ShaderType::VERTEX_SHADER || Ctx.type == ShaderType::HULL_SHADER ||
                                 Ctx.type == ShaderType::DOMAIN_SHADER || Ctx.type == ShaderType::GEOMETRY_SHADER;

    // FIXME: Below heuristics is not a clean design. Revisit this!
    // Start with simd16 or simd32 correspondingly if MinDispatchMode() is 8 or 16, which allows the medium size of
    // space per WI
    //  (simd8: largest, simd32, smallest). In doing so, there will be
    //  some space left for spilling in simd8 if spilling happens.
    int32_t simd_size = isGeometryStageShader
                            ?
                            numLanes(Ctx.platform.getMinDispatchMode())
                            :
                            (Ctx.platform.getMinDispatchMode() == SIMDMode::SIMD8 ? numLanes(SIMDMode::SIMD16)
                                                                                  : numLanes(SIMDMode::SIMD32));
    const int32_t subGrpSize = funcInfoMD->getSubGroupSize()->getSIMDSize();
    if (subGrpSize > simd_size)
      simd_size = std::min(subGrpSize, static_cast<int32_t>(numLanes(SIMDMode::SIMD32)));
    int32_t groupSize = IGCMD::IGCMetaDataHelper::getThreadGroupSize(*pMdUtils, &F);
    if (groupSize == 0)
      groupSize = IGCMD::IGCMetaDataHelper::getThreadGroupSizeHint(*pMdUtils, &F);
    if (groupSize > simd_size)
      simd_size = std::min(groupSize, static_cast<int32_t>(numLanes(SIMDMode::SIMD32)));

    // if one API doesn't support stateless, we should try to use smallest dispatch mode
    // which can hold more pvt_data to avoid error out.
    if (SeparateSpillAndScratch(&Ctx) && !supportsStatelessSpacePrivateMemory)
      simd_size = numLanes(Ctx.platform.getMinDispatchMode());

    unsigned maxScratchSpaceBytes = Ctx.platform.maxPerThreadScratchSpace(
    );
    unsigned scratchSpaceLimitPerWI = maxScratchSpaceBytes / simd_size;
    //
    // If spill happens, since the offset of scratch block rw send message
    // has only 12b, an assertion will be triggered if used scratch space
    // size >= 128 KB, here 128 KB = 2^12 * 256b.
    //
    const unsigned int totalPrivateMemPerWI = getTotalPrivateMemPerWI(&F);

    // FIXME: for now, to shrink size, let's use SIMD8 if have to.
    // later, maybe, we want to change to legacy behavior: SIMD16, to avoid potential spill.
    // but even so, when we support slot0 and slot1, then, we could still use SIMD8.
    if (Ctx.platform.hasScratchSurface() && Ctx.hasSyncRTCalls() && totalPrivateMemPerWI > scratchSpaceLimitPerWI) {
      simd_size = numLanes(Ctx.platform.getMinDispatchMode());
      scratchSpaceLimitPerWI = maxScratchSpaceBytes / simd_size;
    }

    if (totalPrivateMemPerWI > scratchSpaceLimitPerWI) {
      // IGC errors out when we are trying to remove statelesspvtmem of OCL (even though OCl still supports
      // statelesspvtmem). This assertion tests a scenario where (pvt_mem_usage > 256k) while statelessprivatememory is
      // not supported.
      IGC_ASSERT_EXIT(bOCLLegacyStatelessCheck);

      if (!supportsStatelessSpacePrivateMemory) {
        // For XeHP_SDV and above, if any API doesn't support statelesspvtmem, error it out if we find a case where
        // (pvt_mem > 256k). This assertion found a scenario where (pvt_mem_usage > 256k) while statelessprivatememory
        // is not supported.
        IGC_ASSERT(0);
        return true;
      }

      return false;
    }
  }

  // It is safe to use scratch space for private memory.
  return true;
}

unsigned ModuleAllocaAnalysis::getBufferStride(AllocaInst *AI) const {
  IGC_ASSERT(isa<ConstantInt>(AI->getArraySize()));
  Function *F = AI->getParent()->getParent();
  auto allocaInfo = getFuncAllocaInfo(F);
  IGC_ASSERT(allocaInfo);
  auto it = allocaInfo->NonUniformAllocaDesc.find(AI);
  if (it != allocaInfo->NonUniformAllocaDesc.end()) {
    return it->second.second;
  }
  return 0;
}

unsigned ModuleAllocaAnalysis::getBufferOffset(llvm::AllocaInst *AI) const {
  IGC_ASSERT(isa<ConstantInt>(AI->getArraySize()));
  Function *F = AI->getParent()->getParent();
  auto allocaInfo = getFuncAllocaInfo(F);
  IGC_ASSERT(allocaInfo);
  auto uit = allocaInfo->UniformAllocaDesc.find(AI);
  if (uit != allocaInfo->UniformAllocaDesc.end()) {
    return uit->second;
  } else {
    auto nit = allocaInfo->NonUniformAllocaDesc.find(AI);
    IGC_ASSERT(nit != allocaInfo->NonUniformAllocaDesc.end());
    return nit->second.first;
  }
}

unsigned ModuleAllocaAnalysis::isUniform(llvm::AllocaInst *AI) const {
  IGC_ASSERT(isa<ConstantInt>(AI->getArraySize()));
  Function *F = AI->getParent()->getParent();
  auto allocaInfo = getFuncAllocaInfo(F);
  IGC_ASSERT(allocaInfo);
  return allocaInfo->UniformAllocaDesc.count(AI) > 0;
}

llvm::Value *ModuleAllocaAnalysis::getPerThreadOffset(IGCLLVM::IRBuilder<> &IRB, llvm::AllocaInst *AI,
                                                      llvm::Value *simdSize, llvm::Value *threadId,
                                                      bool return64bitOffset /*= false*/) const {
  // Pseudo-code:
  // totalStride = nonUniformBlockOffset + SIMDSize * nonUniformBlockStride
  // totalStride = Align(totalStride, MaxAlignement)
  // perThreadOffset = totalPrivateMemPerWI * HWTID
  Function *F = AI->getParent()->getParent();
  auto allocaInfo = getFuncAllocaInfo(F);
  IGC_ASSERT(allocaInfo);
  Value *nonUniformBlockOffset = IRB.getInt32(allocaInfo->MemoryDescription->GetOffsetForNonUniformAllocas());
  Value *nonUniformBlockStride = IRB.getInt32(allocaInfo->MemoryDescription->GetStrideForNonUniformAllocas());
  Value *nonUniformAllocaBlockSize =
      IRB.CreateMul(nonUniformBlockStride, simdSize, VALUE_NAME(".NonUniformAllocasStride"));
  Value *totalStride =
      IRB.CreateAdd(nonUniformBlockOffset, nonUniformAllocaBlockSize, VALUE_NAME(".TotalPrivateMemPerWI"));
  unsigned maskForTotalPrivateMemPerWI = allocaInfo->MemoryDescription->GetMaxAlignment() - 1u;
  totalStride = IRB.CreateAdd(totalStride, IRB.getInt32(maskForTotalPrivateMemPerWI));
  totalStride =
      IRB.CreateAnd(totalStride, IRB.getInt32(~maskForTotalPrivateMemPerWI), VALUE_NAME("totalPrivateMemStride"));
  if (return64bitOffset) {
    threadId = IRB.CreateZExt(threadId, IRB.getInt64Ty());
    totalStride = IRB.CreateZExt(totalStride, IRB.getInt64Ty());
  }
  Value *perThreadOffset = IRB.CreateMul(threadId, totalStride);
  return perThreadOffset;
}

llvm::Value *ModuleAllocaAnalysis::getOffset(IGCLLVM::IRBuilder<> &IRB, llvm::AllocaInst *AI, llvm::Value *simdSize,
                                             llvm::Value *simdLaneId,
                                             std::optional<uint32_t> perLaneStride /*= std::nullopt*/) const

{
  IGC_ASSERT(isa<ConstantInt>(AI->getArraySize()));
  Function *F = AI->getParent()->getParent();
  auto allocaInfo = getFuncAllocaInfo(F);
  IGC_ASSERT(allocaInfo);

  Value *offset = nullptr;
  auto uit = allocaInfo->UniformAllocaDesc.find(AI);
  if (uit != allocaInfo->UniformAllocaDesc.end()) {
    // Pseudo-code:
    // offset = bufferOffset
    Value *baseOffset = IRB.getInt32(allocaInfo->MemoryDescription->GetOffsetForUniformAllocas());
    Value *relativeOffset = IRB.getInt32(uit->second);
    Value *bufferOffset = IRB.CreateAdd(baseOffset, relativeOffset, VALUE_NAME(AI->getName() + ".BufferOffset"));
    offset = bufferOffset;
  } else {
    // Pseudo-code:
    // offset = nonUniformPrivateMemoryBlockOffset + bufferOffset * SIMDSize + bufferStride * SIMDLane
    auto nit = allocaInfo->NonUniformAllocaDesc.find(AI);
    IGC_ASSERT(nit != allocaInfo->NonUniformAllocaDesc.end());
    auto [localOffset, localStride] = nit->second;
    if (perLaneStride.has_value()) {
      localStride = *perLaneStride;
    }
    Value *localOffsetVal =
        IRB.CreateMul(simdSize, IRB.getInt32(localOffset), VALUE_NAME(AI->getName() + ".SectionOffset"));
    unsigned nonUniformAllocasOffset = allocaInfo->MemoryDescription->GetOffsetForNonUniformAllocas();
    Value *baseOffset = IRB.getInt32(nonUniformAllocasOffset);
    Value *simdBufferOffset = IRB.CreateAdd(baseOffset, localOffsetVal, VALUE_NAME(AI->getName() + ".BufferOffset"));
    bool isUniform = AI->getMetadata("uniform") != nullptr;
    if (!isUniform) {
      Value *localOffset =
          IRB.CreateMul(simdLaneId, IRB.getInt32(localStride), VALUE_NAME(AI->getName() + ".PerLaneOffset"));
      simdBufferOffset = IRB.CreateAdd(simdBufferOffset, localOffset, VALUE_NAME(AI->getName() + ".SIMDBufferOffset"));
    }
    offset = simdBufferOffset;
  }

  return offset;
}

SmallVector<AllocaInst *, 8> &ModuleAllocaAnalysis::getAllocaInsts(Function *F) const {
  return getFuncAllocaInfo(F)->Allocas;
}

unsigned ModuleAllocaAnalysis::getTotalPrivateMemPerWI(Function *F) const {
  // Note that the total size is returned for SIMD1
  // The value is multiplied by the correct SIMD size and
  // reported to UMD. Uniform memory space is shared between
  // all subgroup threads. Therefore, it is safe to divide
  // the value by the smallest simd size.
  // This could be improved if the total size is computed
  // from simd size intrinsic calls instead.
  auto FI = getFuncAllocaInfo(F);
  if (FI) {
    unsigned minSimdSize = getMinSimdSize(F);
    uint64_t totalSize = FI->MemoryDescription->GetTotalSize(minSimdSize);
    totalSize = iSTD::Align(totalSize, minSimdSize);
    totalSize = totalSize / minSimdSize;
    unsigned ret = int_cast<unsigned>(totalSize);

    CodeGenContext &Ctx = *getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
    ModuleMetaData *const modMD = getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
    IGC_ASSERT(nullptr != modMD);
    if (modMD->compOpt.UseScratchSpacePrivateMemory == false &&
        Ctx.m_DriverInfo.supportsStatelessSpacePrivateMemory() &&
        Ctx.m_DriverInfo.requiresPowerOfTwoStatelessSpacePrivateMemorySize()) {
      ret = iSTD::RoundPower2(static_cast<DWORD>(ret));
    }

    return ret;
  }
  return 0;
}

unsigned ModuleAllocaAnalysis::getPrivateMemAlignment(Function *F) const {
  auto FI = getFuncAllocaInfo(F);
  return FI ? FI->MemoryDescription->GetMaxAlignment() : 1;
}

ModuleAllocaAnalysis::FunctionAllocaInfo *ModuleAllocaAnalysis::getFuncAllocaInfo(Function *F) const {
  auto Iter = InfoMap.find(F);
  if (Iter != InfoMap.end())
    return Iter->second;
  return nullptr;
}

ModuleAllocaAnalysis::FunctionAllocaInfo *ModuleAllocaAnalysis::getOrCreateFuncAllocaInfo(Function *F) {
  auto Iter = InfoMap.find(F);
  if (Iter != InfoMap.end())
    return Iter->second;

  auto AllocaInfo = new FunctionAllocaInfo;
  InfoMap[F] = AllocaInfo;
  return AllocaInfo;
}

void ModuleAllocaAnalysis::analyze() {
  IGC_ASSERT(M);
  CodeGenContext &Ctx = *getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
  const IGC::TriboolFlag ForceSupportUniformPrivateMemorySpace =
      static_cast<IGC::TriboolFlag>(IGC_GET_FLAG_VALUE(SupportUniformPrivateMemorySpace));
  switch (ForceSupportUniformPrivateMemorySpace) {
  case IGC::TriboolFlag::Enabled:
    SupportsUniformPrivateMemory = true;
    break;
  case IGC::TriboolFlag::Disabled:
    SupportsUniformPrivateMemory = false;
    break;
  case IGC::TriboolFlag::Default:
    SupportsUniformPrivateMemory = Ctx.m_DriverInfo.supportsUniformPrivateMemorySpace();
    break;
  }

  if (FGA && FGA->getModule()) {
    IGC_ASSERT(FGA->getModule() == M);
    for (auto FG : *FGA)
      analyze(FG);
  } else {
    for (auto &F : M->getFunctionList()) {
      if (F.empty())
        continue;

      PrivateMemoryDescription &MemoryDescription = MemoryDescriptions.emplace_back();
      analyze(&F, MemoryDescription);
    }
  }
}

void ModuleAllocaAnalysis::analyze(IGC::FunctionGroup *FG) {
  // Calculate the size of private-memory we need to allocate to
  // every function-sub-group. Eache sub-group is led by a kernel or
  // a stack-call function.
  // Note that the function order does affect the final total amount of
  // private memory due to possible alignment constraints.
  for (auto SubG : FG->Functions) {
    PrivateMemoryDescription &MemoryDescription = MemoryDescriptions.emplace_back();
    for (Function *F : *SubG) {
      if (F->empty())
        continue;
      analyze(F, MemoryDescription);
    }
  }
}

void ModuleAllocaAnalysis::analyze(Function *F, PrivateMemoryDescription &MemoryDescription) {
  const DataLayout *DL = &M->getDataLayout();

  // Create alloca info even when there is no alloca, so that each function gets
  // an info entry.
  FunctionAllocaInfo *AllocaInfo = getOrCreateFuncAllocaInfo(F);
  AllocaInfo->MemoryDescription = &MemoryDescription;

  // Collect allocas.
  SmallVector<AllocaInst *, 8> Allocas;
  for (auto &BB : *F) {
    for (auto &Inst : BB) {
      if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
        Allocas.push_back(AI);
      }
    }
  }

  if (Allocas.empty())
    return;

  std::sort(Allocas.begin(), Allocas.end(), [&DL](AllocaInst *AI1, AllocaInst *AI2) {
    return FunctionAllocaInfo::getAlignment(AI1, DL) < FunctionAllocaInfo::getAlignment(AI2, DL);
  });

  for (auto AI : Allocas) {
    AllocaInfo->AssignAlloca(AI, DL, SupportsUniformPrivateMemory);
  }

  // Update collected allocas into the function alloca info object.
  AllocaInfo->Allocas.swap(Allocas);
}

unsigned ModuleAllocaAnalysis::FunctionAllocaInfo::getAlignment(llvm::AllocaInst *AI, const llvm::DataLayout *DL) {
  alignment_t Alignment = IGCLLVM::getAlignmentValue(AI);
  if (Alignment == 0)
    Alignment = DL->getABITypeAlign(AI->getAllocatedType()).value();
  return std::max(int_cast<unsigned>(Alignment), 1u);
}

unsigned ModuleAllocaAnalysis::FunctionAllocaInfo::getSize(llvm::AllocaInst *AI, const llvm::DataLayout *DL) {
  IGC_ASSERT(isa<ConstantInt>(AI->getArraySize()));
  ConstantInt *const SizeVal = cast<ConstantInt>(AI->getArraySize());
  IGC_ASSERT(nullptr != SizeVal);
  unsigned CurSize = (unsigned)(SizeVal->getZExtValue() * DL->getTypeAllocSize(AI->getAllocatedType()));
  bool isPackedStruct = false;
  if (auto st = dyn_cast<StructType>(AI->getAllocatedType())) {
    isPackedStruct = st->isPacked();
  }
  if (!isPackedStruct) {
    auto Alignment = getAlignment(AI, DL);
    CurSize = iSTD::Align(CurSize, Alignment);
  }
  return CurSize;
}

void ModuleAllocaAnalysis::FunctionAllocaInfo::AssignAlloca(llvm::AllocaInst *AI, const llvm::DataLayout *DL,
                                                            bool SupportsUniformPrivateMemory) {
  // Dynamic allocas are handled by stack-related instructions
  if (!isa<ConstantInt>(AI->getArraySize())) {
    return;
  }
  bool isUniform = AI->getMetadata("uniform") != nullptr;
  if (SupportsUniformPrivateMemory && isUniform) {
    unsigned alignement = getAlignment(AI, DL);
    MemoryDescription->MaximumAlignmentForUniformAllocas =
        std::max(MemoryDescription->MaximumAlignmentForUniformAllocas, alignement);
    MemoryDescription->TotalSizeForUniformAllocas =
        iSTD::Align(MemoryDescription->TotalSizeForUniformAllocas, alignement);
    UniformAllocaDesc[AI] = MemoryDescription->TotalSizeForUniformAllocas;
    MemoryDescription->TotalSizeForUniformAllocas += getSize(AI, DL);
  } else {
    unsigned alignement = getAlignment(AI, DL);
    MemoryDescription->MaximumAlignmentForNonUniformAllocas =
        std::max(MemoryDescription->MaximumAlignmentForNonUniformAllocas, alignement);
    unsigned stride = getSize(AI, DL);
    MemoryDescription->TotalStrideForNonUniformAllocas =
        iSTD::Align(MemoryDescription->TotalStrideForNonUniformAllocas, alignement);
    NonUniformAllocaDesc[AI] = std::make_pair(MemoryDescription->TotalStrideForNonUniformAllocas, stride);
    MemoryDescription->TotalStrideForNonUniformAllocas += stride;
  }
}

unsigned ModuleAllocaAnalysis::getMinSimdSize(llvm::Function *pFunc) const {
  ModuleMetaData &modMD = *getAnalysis<MetaDataUtilsWrapper>().getModuleMetaData();
  CodeGenContext &Ctx = *getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
  uint32_t minSimdSize = numLanes(Ctx.platform.getMinDispatchMode());
  switch (Ctx.type) {
  case ShaderType::VERTEX_SHADER:
  case ShaderType::HULL_SHADER:
  case ShaderType::DOMAIN_SHADER:
  case ShaderType::GEOMETRY_SHADER:
    minSimdSize = numLanes(Ctx.GetSIMDMode());
    break;
  case ShaderType::TASK_SHADER:
    minSimdSize = modMD.taskInfo.SubgroupSize > 0 ? modMD.taskInfo.SubgroupSize : minSimdSize;
    break;
  case ShaderType::MESH_SHADER:
    minSimdSize = modMD.msInfo.SubgroupSize > 0 ? modMD.msInfo.SubgroupSize : minSimdSize;
    break;
  case ShaderType::COMPUTE_SHADER:
    minSimdSize = modMD.csInfo.waveSize > 0 ? modMD.csInfo.waveSize : minSimdSize;
    minSimdSize = modMD.csInfo.forcedSIMDSize > 0 ? modMD.csInfo.forcedSIMDSize : minSimdSize;
    break;
  case ShaderType::PIXEL_SHADER:
    minSimdSize = modMD.compOpt.forcePixelShaderSIMDMode > 0
                      ? BIT(iSTD::bsf(modMD.compOpt.forcePixelShaderSIMDMode)) << 3
                      : minSimdSize;
    break;
  case ShaderType::OPENCL_SHADER: {
    IGCMD::MetaDataUtils *pMdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
    auto funcInfoMD = pMdUtils->getFunctionsInfoItem(pFunc);
    const int32_t subGrpSize = funcInfoMD->getSubGroupSize()->getSIMDSize();
    if (subGrpSize > 0 && int_cast<uint16_t>(subGrpSize) > int_cast<uint16_t>(minSimdSize))
      minSimdSize = std::min(int_cast<uint16_t>(subGrpSize), numLanes(SIMDMode::SIMD32));
    int32_t groupSize = IGCMD::IGCMetaDataHelper::getThreadGroupSize(*pMdUtils, pFunc);
    if (groupSize == 0)
      groupSize = IGCMD::IGCMetaDataHelper::getThreadGroupSizeHint(*pMdUtils, pFunc);
    if (groupSize > 0 && int_cast<uint16_t>(groupSize) > int_cast<uint16_t>(minSimdSize))
      minSimdSize = std::min(int_cast<uint16_t>(groupSize), numLanes(SIMDMode::SIMD32));
    break;
  }
  default:
    break;
  }
  return minSimdSize;
}

char ModuleAllocaAnalysis::ID = 0;

// Register pass to igc-opt
#define PASS_FLAG "igc-module-alloca-info"
#define PASS_DESCRIPTION "Analyse memory usage based on alloca instructions"
#define PASS_CFG_ONLY true
#define PASS_ANALYSIS true
IGC_INITIALIZE_PASS_BEGIN(ModuleAllocaAnalysis, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_END(ModuleAllocaAnalysis, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

} // namespace IGC