Improve InlineHelper LLVM utility

Improve InlineHelper LLVM utility

Author: jaladreips

IGC/common/LLVMUtils.cpp

@@ -951,33 +951,47 @@ void InlineHelper::InlineAndOptimize(CallInst* callInst)
         m_calledFunction = callInst->getCalledFunction();
 #endif
     auto* fn = callInst->getFunction();
-
     InlineFunctionInfo IFI;
     bool CanInline = IGCLLVM::InlineFunction(*callInst, IFI);
     IGC_ASSERT_MESSAGE(CanInline, "failed to inline?");
 
-    auto& perFnAllocas = m_InlinedStaticArrayAllocas[fn];
-
     // Merge static array allocas to reduce the use of private
     // memory. This is a similar optimization that exists in
     // the inliner, see mergeInlinedArrayAllocas().
     llvm::erase_if(
         IFI.StaticAllocas,
         [](AllocaInst* I)
         {
-            return I->isArrayAllocation() || !I->getAllocatedType()->isArrayTy();
+            return I->getAllocatedType()->isArrayTy() || !isa<ConstantInt>(I->getArraySize());
         }
     );
 
-    auto staticAlloca = IFI.StaticAllocas.begin();
-    for (auto* fnAlloca : perFnAllocas)
+    AllocaMap localAllocas;
+    for (auto* I : IFI.StaticAllocas)
     {
-        if ((*staticAlloca)->getAllocatedType() == fnAlloca->getAllocatedType())
-            (*staticAlloca)->replaceAllUsesWith(fnAlloca);
-
-        if (++staticAlloca == IFI.StaticAllocas.end())
-            return;
+        localAllocas[
+            std::make_tuple(
+                I->getType(),
+                cast<ConstantInt>(I->getArraySize())->getZExtValue(),
+                I->getAddressSpace())
+        ].push_back(I);
     }
 
-    perFnAllocas.append(staticAlloca, IFI.StaticAllocas.end());
+    for (auto& [key, perTypeAllocas] :localAllocas)
+    {
+        auto localAlloca = perTypeAllocas.begin();
+        auto& globalAllocas = m_InlinedStaticArrayAllocas[fn][key];
+        for (auto globalAlloca : globalAllocas)
+        {
+            (*localAlloca)->replaceAllUsesWith(globalAlloca);
+            (*localAlloca)->eraseFromParent();
+
+            localAlloca++;
+
+            if (localAlloca == perTypeAllocas.end())
+                break;
+        }
+
+        globalAllocas.append(localAlloca, perTypeAllocas.end());
+    }
 }

IGC/common/LLVMUtils.h

@@ -52,7 +52,14 @@ void DumpHashToOptions(const ShaderHash&, const ShaderType);
 class InlineHelper
 {
 private:
-    llvm::DenseMap<llvm::Function*, llvm::SmallVector<llvm::AllocaInst*, 4>> m_InlinedStaticArrayAllocas;
+    // create a map of (type, array size, address space to alloca instructions)
+    using AllocaMap = llvm::DenseMap<
+        std::tuple<llvm::Type*, uint64_t, uint32_t>,
+        llvm::SmallVector<llvm::AllocaInst*>
+    >;
+
+    llvm::DenseMap<llvm::Function*, AllocaMap> m_InlinedStaticArrayAllocas;
+
 #if defined(_RELEASE_INTERNAL) || defined(_DEBUG)
     llvm::Function* m_calledFunction = nullptr;
 #endif