Make MergeAllocas handle different sized arrays

Merge differently sized non-overlapping array allocas.

Author: stefan-il

IGC/AdaptorCommon/RayTracing/MergeAllocas.cpp

@@ -18,6 +18,7 @@ SPDX-License-Identifier: MIT
 #include <llvm/Analysis/LoopInfo.h>
 #include <llvm/IR/Constants.h>
 #include <llvm/IR/Dominators.h>
+#include <llvm/IR/IRBuilder.h>
 #include <llvm/IR/InstIterator.h>
 #include <llvm/IR/Instructions.h>
 #include "common/LLVMWarningsPop.hpp"
@@ -332,7 +333,7 @@ void MergeAllocas::getAnalysisUsage(llvm::AnalysisUsage& AU) const
 
 bool MergeAllocas::runOnFunction(Function& F)
 {
-    if (F.hasOptNone()){
+    if (skipFunction(F)){
         return false;
     }
 
@@ -341,7 +342,7 @@ bool MergeAllocas::runOnFunction(Function& F)
     // we group the allocations by type, then sort them into buckets with nonoverlapping liveranges
     // can this be generalized into allocas for types of the same size, not only types?
     using BucketT = SmallVector<std::pair<Instruction*, AllocationBasedLivenessAnalysis::LivenessData*>>;
-    DenseMap<std::tuple<llvm::Type*, uint64_t, uint32_t>, SmallVector<BucketT>> buckets;
+    DenseMap<std::tuple<llvm::Type*, uint32_t, uint32_t>, SmallVector<BucketT>> buckets;
 
     for (const auto& A : ABLA)
     {
@@ -353,9 +354,11 @@ bool MergeAllocas::runOnFunction(Function& F)
         if (!isa<ConstantInt>(AI->getArraySize()))
             continue;
 
+        auto* AllocatedType = AI->getAllocatedType();
+        uint32_t IsArray = AllocatedType->isArrayTy()? 1 : 0;
         auto& perTypeBuckets = buckets[std::make_tuple(
-            AI->getAllocatedType(),
-            cast<ConstantInt>(AI->getArraySize())->getZExtValue(),
+            (IsArray)? AllocatedType->getArrayElementType() : AllocatedType,
+            IsArray,
             AI->getAddressSpace()
         )];
 
@@ -379,7 +382,7 @@ bool MergeAllocas::runOnFunction(Function& F)
 
     bool changed = false;
 
-    for (const auto& [_, perTypeBuckets] : buckets)
+    for (const auto& [allocaType, perTypeBuckets] : buckets)
     {
         for (const auto& bucket : perTypeBuckets)
         {
@@ -388,21 +391,36 @@ bool MergeAllocas::runOnFunction(Function& F)
                 continue;
             }
 
+            bool IsArray = std::get<1>(allocaType);
             Instruction* firstAlloca = nullptr;
+            if (IsArray)
+            {
+                firstAlloca = std::max_element(bucket.begin(), bucket.end(), [](const auto& a, const auto& b) {
+                    return cast<AllocaInst>(a.first)->getAllocatedType()->getArrayNumElements() < cast<AllocaInst>(b.first)->getAllocatedType()->getArrayNumElements();
+                })->first;
+            }
+            else
+            {
+                firstAlloca = bucket[0].first;
+            }
+            firstAlloca->moveBefore(F.getEntryBlock().getFirstNonPHI());
+            firstAlloca->setName(VALUE_NAME("MergedAlloca"));
             for (const auto& [I, _] : bucket)
             {
-                if (!firstAlloca)
+                if (firstAlloca == I)
                 {
-                    firstAlloca = I;
-                    firstAlloca->moveBefore(F.getEntryBlock().getFirstNonPHI());
-                    firstAlloca->setName(VALUE_NAME("MergedAlloca"));
+                    continue;
                 }
-                else
+                auto* ReplacementValue = firstAlloca;
+                if (firstAlloca->getType() != I->getType())
                 {
-                    I->replaceAllUsesWith(firstAlloca);
-                    I->eraseFromParent();
-                }
-
+                    IRBuilder<> Builder(firstAlloca->getParent());
+                    Builder.SetInsertPoint(firstAlloca->getNextNode());
+                    auto *CastedValue = llvm::cast<Instruction>(Builder.CreateBitCast(firstAlloca, I->getType()));
+                    ReplacementValue = CastedValue;
+                };
+                I->replaceAllUsesWith(ReplacementValue);
+                I->eraseFromParent();
             }
 
             changed = true;

IGC/Compiler/tests/PrivateMemoryResolution/MergeAllocas/merge_allocas_arrays_and_regular.ll

@@ -0,0 +1,37 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+; RUN: igc_opt --igc-merge-allocas --igc-private-mem-resolution -S %s --platformpvc | FileCheck %s
+; ------------------------------------------------
+; PrivateMemoryResolution
+; ------------------------------------------------
+
+; Check that array allocas are not merged with regular allocas.
+
+define spir_kernel void @main(float addrspace(1)* %0, i64 %1, i64 %2, i32 %3, i32 %4) {
+; CHECK-LABEL: main
+; CHECK-NEXT: alloca [128 x float], align 4
+; CHECK-NEXT: alloca float, align 4
+  %6 = alloca [128 x float], align 4
+  %7 = alloca float, align 4
+  br label %8
+
+8:
+  %9 = icmp ult i32 %3, %4
+  br i1 %9, label %10, label %12
+
+10:
+  %11 = getelementptr inbounds [128 x float], [128 x float]* %6, i64 0, i64 %1
+  br label %14
+
+12:
+  %13 = load float, float* %7, align 4
+  br label %14
+
+14:
+  ret void
+}

IGC/Compiler/tests/PrivateMemoryResolution/MergeAllocas/merge_allocas_different_sized_arrays.ll

@@ -0,0 +1,38 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+; RUN: igc_opt --typed-pointers --igc-merge-allocas --igc-private-mem-resolution -S %s --platformpvc | FileCheck %s
+; ------------------------------------------------
+; PrivateMemoryResolution
+; ------------------------------------------------
+
+; Check that allocas are merged and that only largest array is kept.
+
+define spir_kernel void @main(float addrspace(1)* %0, i64 %1, i64 %2, i32 %3, i32 %4) {
+; CHECK-LABEL: main
+; CHECK-NEXT: alloca [256 x float], align 4
+; CHECK-NEXT: bitcast [256 x float]* {{.*}} to [128 x float]*
+; CHECK-NOT: alloca [128 x float], align 4
+  %6 = alloca [128 x float], align 4
+  %7 = alloca [256 x float], align 4
+  br label %8
+
+8:
+  %9 = icmp ult i32 %3, %4
+  br i1 %9, label %10, label %12
+
+10:
+  %11 = getelementptr inbounds [128 x float], [128 x float]* %6, i64 0, i64 %1
+  br label %14
+
+12:
+  %13 = getelementptr inbounds [256 x float], [256 x float]* %7, i64 0, i64 %2
+  br label %14
+
+14:
+  ret void
+}

IGC/Compiler/tests/PrivateMemoryResolution/MergeAllocas/merge_allocas_different_type_arrays.ll

@@ -0,0 +1,37 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+; RUN: igc_opt --typed-pointers --igc-merge-allocas --igc-private-mem-resolution -S %s --platformpvc | FileCheck %s
+; ------------------------------------------------
+; PrivateMemoryResolution
+; ------------------------------------------------
+
+; Check that allocas are not merged if array elements have different types.
+
+define spir_kernel void @main(float addrspace(1)* %0, i64 %1, i64 %2, i32 %3, i32 %4) {
+; CHECK-LABEL: main
+; CHECK-NEXT: alloca [128 x float], align 4
+; CHECK-NEXT: alloca [256 x i64], align 4
+  %6 = alloca [128 x float], align 4
+  %7 = alloca [256 x i64], align 4
+  br label %8
+
+8:
+  %9 = icmp ult i32 %3, %4
+  br i1 %9, label %10, label %12
+
+10:
+  %11 = getelementptr inbounds [128 x float], [128 x float]* %6, i64 0, i64 %1
+  br label %14
+
+12:
+  %13 = getelementptr inbounds [256 x i64], [256 x i64]* %7, i64 0, i64 %2
+  br label %14
+
+14:
+  ret void
+}