Handle loops in liveness analysis for merge allocas

Liveness analysis does not consider the declaration (allocation instruction) as a part of live range of the allocation, instead finds the last possible point that still dominates all uses.
That did not work with loops, where the analysis did not consider the need for the live range to wrap around the loop.
The fix introduced by this change checks if the live range escapes the loop (flows out the exiting block and into the exit block), and if that's true, includes the whole loop in the live range.

Author: jaladreips

IGC/AdaptorCommon/RayTracing/MergeAllocas.cpp

@@ -14,10 +14,11 @@ SPDX-License-Identifier: MIT
 #include "common/LLVMWarningsPush.hpp"
 #include <llvm/ADT/SetVector.h>
 #include <llvm/ADT/SetOperations.h>
+#include <llvm/Analysis/LoopInfo.h>
+#include <llvm/IR/Constants.h>
 #include <llvm/IR/Dominators.h>
 #include <llvm/IR/InstIterator.h>
 #include <llvm/IR/Instructions.h>
-#include <llvm/IR/Constants.h>
 #include "common/LLVMWarningsPop.hpp"
 
 using namespace llvm;
@@ -26,6 +27,7 @@ using namespace IGC;
 // Register pass to igc-opt
 IGC_INITIALIZE_PASS_BEGIN(AllocationBasedLivenessAnalysis, "igc-allocation-based-liveness-analysis", "Analyze the lifetimes of instruction allocated by a specific intrinsic", false, true)
 IGC_INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 IGC_INITIALIZE_PASS_END(AllocationBasedLivenessAnalysis, "igc-allocation-based-liveness-analysis", "Analyze the lifetimes of instruction allocated by a specific intrinsic", false, true)
 
 char AllocationBasedLivenessAnalysis::ID = 0;
@@ -34,6 +36,7 @@ void AllocationBasedLivenessAnalysis::getAnalysisUsage(llvm::AnalysisUsage& AU)
 {
     AU.setPreservesAll();
     AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<LoopInfoWrapperPass>();
 }
 
 AllocationBasedLivenessAnalysis::AllocationBasedLivenessAnalysis() : FunctionPass(ID)
@@ -68,6 +71,7 @@ AllocationBasedLivenessAnalysis::LivenessData* AllocationBasedLivenessAnalysis::
     // that's a practice, but we only care about the last block that dominates all uses
     BasicBlock* commonDominator = nullptr;
     auto* DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+    auto* LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
 
     bool hasNoLifetimeEnd = false;
 
@@ -142,10 +146,10 @@ AllocationBasedLivenessAnalysis::LivenessData* AllocationBasedLivenessAnalysis::
         );
     }
 
-    return new LivenessData(I, allUsers, commonDominator);
+    return new LivenessData(I, allUsers, *LI, commonDominator);
 }
 
-AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocationInstruction, SetVector<Instruction*> usersOfAllocation, BasicBlock* userDominatorBlock)
+AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocationInstruction, const SetVector<Instruction*>& usersOfAllocation, const LoopInfo& LI, BasicBlock* userDominatorBlock)
 {
     if (!userDominatorBlock)
         userDominatorBlock = allocationInstruction->getParent();
@@ -176,13 +180,39 @@ AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocat
         }
     }
 
+    // if the lifetime escapes any loop, we should make sure all the loops blocks are included
+    for (const auto& loop : LI)
+    {
+        SmallVector<std::pair<BasicBlock*, BasicBlock*>> exitEdges;
+        loop->getExitEdges(exitEdges);
+
+        if (llvm::any_of(exitEdges, [&](auto edge) { return bbOut.contains(edge.first) && bbIn.contains(edge.second); }))
+        {
+            llvm::for_each(
+                loop->blocks(),
+                [&](auto* block) {
+                    bbOut.insert(block);
+                    if (block != loop->getHeader())
+                        bbIn.insert(block);
+                }
+            );
+        }
+    }
+
+    // substract the inflow blocks from the outflow blocks to find the block which starts the lifetime - there should be only one!
+    auto bbOutOnly = bbOut;
+    set_subtract(bbOutOnly, bbIn);
+
+    IGC_ASSERT_MESSAGE(bbOutOnly.size() == 1, "Multiple lifetime start blocks?");
+
+    auto* lifetimeStartBB = *bbOutOnly.begin();
+
     // fill out the lifetime start/ends instruction
-    for (auto& I : *userDominatorBlock)
+    for (auto& I : *lifetimeStartBB)
     {
         lifetimeStart = &I;
         if (usersOfAllocation.contains(&I))
             break;
-
     }
 
     // if bbIn is empty, the entire lifetime is contained within userDominatorBlock

IGC/AdaptorCommon/RayTracing/MergeAllocas.h

@@ -7,11 +7,18 @@ SPDX-License-Identifier: MIT
 ============================= end_copyright_notice ===========================*/
 
 #include "common/LLVMWarningsPush.hpp"
-#include <llvm/IR/Instruction.h>
+#include <llvm/Pass.h>
+#include <llvm/ADT/DenseSet.h>
 #include <llvm/ADT/SetVector.h>
-#include <llvm/ADT/MapVector.h>
 #include "common/LLVMWarningsPop.hpp"
 
+namespace llvm {
+    class BasicBlock;
+    class Function;
+    class Instruction;
+    class LoopInfo;
+}
+
 namespace IGC
 {
     class MergeAllocas : public llvm::FunctionPass
@@ -58,7 +65,7 @@ namespace IGC
             llvm::DenseSet<llvm::BasicBlock*> bbIn;
             llvm::DenseSet<llvm::BasicBlock*> bbOut;
 
-            LivenessData(llvm::Instruction* allocationInstruction, llvm::SetVector<llvm::Instruction*> usersOfAllocation, llvm::BasicBlock* userDominatorBlock = nullptr);
+            LivenessData(llvm::Instruction* allocationInstruction, const llvm::SetVector<llvm::Instruction*>& usersOfAllocation, const llvm::LoopInfo& LI, llvm::BasicBlock* userDominatorBlock = nullptr);
 
             bool OverlapsWith(const LivenessData& LD) const;
         };
@@ -78,5 +85,4 @@ namespace IGC
             }
         }
     };
-
 } // namespace IGC

IGC/Compiler/tests/MergeAllocas/test1.ll renamed to IGC/Compiler/tests/MergeAllocas/basic-1.ll

@@ -1,3 +1,11 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
 target datalayout = "e-p:32:32:32-p1:64:64:64-p2:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:32:32-v96:32:32-v128:32:32-a0:0:32-n8:16:32-S32"
 target triple = "dxil-ms-dx"
 

IGC/Compiler/tests/MergeAllocas/test2.ll renamed to IGC/Compiler/tests/MergeAllocas/basic-2.ll

@@ -1,3 +1,11 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
 target datalayout = "e-p:32:32:32-p1:64:64:64-p2:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:32:32-v96:32:32-v128:32:32-a0:0:32-n8:16:32-S32"
 target triple = "dxil-ms-dx"
 

IGC/Compiler/tests/MergeAllocas/test3.ll renamed to IGC/Compiler/tests/MergeAllocas/basic-3.ll

@@ -1,3 +1,11 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
 target datalayout = "e-p:32:32:32-p1:64:64:64-p2:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:32:32-v96:32:32-v128:32:32-a0:0:32-n8:16:32-S32"
 target triple = "dxil-ms-dx"
 

IGC/Compiler/tests/MergeAllocas/test4.ll renamed to IGC/Compiler/tests/MergeAllocas/basic-4.ll

@@ -1,3 +1,11 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
 target datalayout = "e-p:32:32:32-p1:64:64:64-p2:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:32:32-v96:32:32-v128:32:32-a0:0:32-n8:16:32-S32"
 target triple = "dxil-ms-dx"
 

IGC/Compiler/tests/MergeAllocas/test5.ll renamed to IGC/Compiler/tests/MergeAllocas/basic-5.ll

@@ -1,3 +1,11 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
 target datalayout = "e-p:32:32:32-p1:64:64:64-p2:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:32:32-v96:32:32-v128:32:32-a0:0:32-n8:16:32-S32"
 target triple = "dxil-ms-dx"
 

IGC/Compiler/tests/MergeAllocas/loop-1.ll

@@ -0,0 +1,56 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
+target datalayout = "e-p:32:32:32-p1:64:64:64-p2:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:32:32-v96:32:32-v128:32:32-a0:0:32-n8:16:32-S32"
+target triple = "dxil-ms-dx"
+
+; RUN: igc_opt --opaque-pointers -S --platformbmg --igc-merge-allocas %s | FileCheck %s
+
+define void @test() #0 {
+  %alloca1 = alloca i32
+  %alloca2 = alloca i32
+; CHECK: alloca i32
+; CHECK: alloca i32
+  br label %looppreheader
+
+looppreheader:
+
+  br label %loopheader
+
+loopheader:
+  br label %loopbody1
+
+loopbody1:
+  store i32 9, ptr %alloca1
+  br label %loopbody2
+
+loopbody2:
+  %a = load i32, ptr %alloca1
+  store i32 10, ptr %alloca2
+  br i1 false, label %loopheader, label %loopexit
+
+loopexit:
+  %load1 = load i32, ptr %alloca2
+  br label %commonret
+
+commonret:
+  ret void
+}
+
+attributes #0 = { nounwind readnone }
+
+
+!llvm.ident = !{!0, !0, !0, !0}
+!dx.version = !{!1, !1, !1, !1}
+!dx.valver = !{!2, !2, !2, !2}
+!dx.shaderModel = !{!3, !3, !3, !3}
+
+!0 = !{!"clang version 3.7 (tags/RELEASE_370/final)"}
+!1 = !{i32 1, i32 5}
+!2 = !{i32 1, i32 6}
+!3 = !{!"lib", i32 6, i32 5}