Fixes for allocation based liveness analysis

Fixes:
 * The assumption a loop header is a good place to start the lifetime is incorrect. Loop header is a part of the loop, so we should either choose the preheader or, if there is none, use the domtree to find common dominator for all predecessors.
 * Infinite lifetime handling would color the entire function if the first use of the allocation would not dominate all return instructions. This has been changed to a simple reachability search algorithm that will find all the blocks reachable from the liveness end points. It's probably too much, but it will not go back all the way to the functions entry point.
 * Only mark basic block as inflowing or outflowing if they actually have predecessors or successors. Previously, a function entry point would be marked as "inflow" even though there aren't any blocks pointing to it.
(cherry picked from commit 40d5e01)

Author: jaladreips

IGC/AdaptorCommon/RayTracing/MergeAllocas.cpp

@@ -95,7 +95,7 @@ AllocationBasedLivenessAnalysis::LivenessData* AllocationBasedLivenessAnalysis::
     }
 
     // figure out the potential accesses to the memory via GEP and bitcasts
-    while (!worklist.empty() && !hasNoLifetimeEnd)
+    while (!worklist.empty())
     {
         auto* use = worklist.pop_back_val();
         auto* II = cast<Instruction>(use->getUser());
@@ -139,22 +139,44 @@ AllocationBasedLivenessAnalysis::LivenessData* AllocationBasedLivenessAnalysis::
         }
     }
 
-    // we add the return instructions to the list of users to express the infinite lifetime
-    if (hasNoLifetimeEnd)
+    return new LivenessData(I, allUsers, *LI, *DT, commonDominator, hasNoLifetimeEnd);
+}
+
+template<typename range>
+static inline void doWorkLoop(
+    SmallVector<BasicBlock*>& worklist,
+    DenseSet<BasicBlock*>& bbSet1,
+    DenseSet<BasicBlock*>& bbSet2,
+    std::function<range(BasicBlock*)> iterate,
+    std::function<bool(BasicBlock*)> continueCondition
+) {
+    // perform data flow analysis
+    while (!worklist.empty())
     {
-        for_each(instructions(*I->getFunction()),
-            [&](auto& II)
-            {
-                if (isa<ReturnInst>(&II))
-                    allUsers.insert(&II);
-            }
-        );
+        auto* currbb = worklist.pop_back_val();
+
+        if (continueCondition(currbb))
+            continue;
+
+        bool addToSet1 = false;
+
+        for (auto* pbb : iterate(currbb))
+        {
+            addToSet1 = true;
+
+            bool inserted = bbSet2.insert(pbb).second;
+
+            if (inserted)
+                worklist.push_back(pbb);
+        }
+
+        if (addToSet1)
+            bbSet1.insert(currbb);
     }
 
-    return new LivenessData(I, allUsers, *LI, commonDominator);
 }
 
-AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocationInstruction, const SetVector<Instruction*>& usersOfAllocation, const LoopInfo& LI, BasicBlock* userDominatorBlock)
+AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocationInstruction, const SetVector<Instruction*>& usersOfAllocation, const LoopInfo& LI, const DominatorTree& DT, BasicBlock* userDominatorBlock, bool isLifetimeInfinite)
 {
     if (!userDominatorBlock)
         userDominatorBlock = allocationInstruction->getParent();
@@ -171,38 +193,44 @@ AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocat
     // Keep track of loop header of blocks that contain allocation instruction
     auto* allocationParent = allocationInstruction->getParent();
     llvm::SmallPtrSet<llvm::BasicBlock*, 4> containedLoopHeaders;
-    if (const auto* parentLoop = LI.getLoopFor(allocationParent);
-        parentLoop != nullptr) {
+    if (const auto* parentLoop = LI.getLoopFor(allocationParent))
+    {
         containedLoopHeaders.insert(parentLoop->getHeader());
-        while (parentLoop->getParentLoop() != nullptr) {
+        while (parentLoop->getParentLoop()) {
             parentLoop = parentLoop->getParentLoop();
             containedLoopHeaders.insert(parentLoop->getHeader());
         }
     }
+
     // perform data flow analysis
-    while (!worklist.empty())
+    doWorkLoop<llvm::pred_range>(
+        worklist,
+        bbIn,
+        bbOut,
+        [&](auto* currbb) { return llvm::predecessors(currbb); },
+        [&](auto* currbb) { return bbIn.contains(currbb) || currbb == userDominatorBlock || containedLoopHeaders.contains(currbb); }
+    );
+
+    // handle infinite lifetime
+    if (isLifetimeInfinite)
     {
-        auto* currbb = worklist.pop_back_val();
-
-        if (bbIn.contains(currbb) || currbb == userDominatorBlock)
-            continue;
-
-        // If alloca is defined in the loop, we skip loop header
-        // so that we don't escape loop scope.
-        if (containedLoopHeaders.count(currbb) != 0)
-        {
-            continue;
-        }
-
-        if (currbb != allocationParent)
-        {
-            bbIn.insert(currbb);
-        }
-        for (auto* pbb : llvm::predecessors(currbb))
-        {
-            bbOut.insert(pbb);
-            worklist.push_back(pbb);
-        }
+        // traverse all the successors until there are no left.
+        auto bbInOnly = bbIn;
+        set_subtract(bbInOnly, bbOut);
+
+        for (auto* bb : bbInOnly)
+            worklist.push_back(bb);
+
+        // in case the only use is the one that causes lifetime escape
+        worklist.push_back(userDominatorBlock);
+
+        doWorkLoop<llvm::succ_range>(
+            worklist,
+            bbOut,
+            bbIn,
+            [&](auto* currbb) { return llvm::successors(currbb); },
+            [&](auto* currbb) { return false; }
+        );
     }
 
     // if the lifetime escapes any loop, we should make sure all the loops blocks are included
@@ -215,12 +243,38 @@ AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocat
         {
             llvm::for_each(
                 loop->blocks(),
-                [&](auto* block) {
-                    bbOut.insert(block);
-                    if (block != loop->getHeader())
-                        bbIn.insert(block);
-                }
+                [&](auto* block) { bbOut.insert(block); bbIn.insert(block); }
             );
+
+            if (loop->getLoopPreheader())
+            {
+                bbOut.insert(loop->getLoopPreheader());
+            }
+            else
+            {
+                // if the header has multiple predecessors, we need to find the common dominator of all of these
+                auto* commonDominator = loop->getHeader();
+                for (auto* bb : llvm::predecessors(loop->getHeader()))
+                {
+                    if (loop->contains(bb))
+                        continue;
+
+                    commonDominator = DT.findNearestCommonDominator(commonDominator, bb);
+                    worklist.push_back(bb);
+                }
+
+                // acknowledge lifetime flow out of the common dominator block
+                bbOut.insert(commonDominator);
+
+                // add all blocks inbetween
+                doWorkLoop<llvm::pred_range>(
+                    worklist,
+                    bbIn,
+                    bbOut,
+                    [&](auto* currbb) { return llvm::predecessors(currbb); },
+                    [&](auto* currbb) { return bbOut.contains(currbb) || currbb == commonDominator; }
+                );
+            }
         }
     }
 
@@ -264,7 +318,7 @@ AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocat
         {
             for (auto& I : llvm::reverse(*bb))
             {
-                if (usersOfAllocation.contains(&I))
+                if (usersOfAllocation.contains(&I) || isLifetimeInfinite)
                 {
                     lifetimeEnds.push_back(&I);
                     break;
@@ -289,12 +343,34 @@ bool AllocationBasedLivenessAnalysis::LivenessData::OverlapsWith(const LivenessD
     // check lifetime boundaries
     for (auto& [LD1, LD2] : { std::make_pair(*this, LD), std::make_pair(LD, *this) })
     {
-        if (LD1.lifetimeEnds.size() == 1 && *LD1.lifetimeEnds.begin() == LD1.lifetimeStart)
-            continue;
-
         for (auto* I : LD1.lifetimeEnds)
         {
-            if (I->getParent() == LD2.lifetimeStart->getParent())
+            // what if LD1 is contained in a single block
+            if (I->getParent() == LD1.lifetimeStart->getParent())
+            {
+                auto* bb = I->getParent();
+                bool inflow = LD2.bbIn.contains(bb);
+                bool outflow = LD2.bbOut.contains(bb);
+                bool lifetimeStart = LD2.lifetimeStart->getParent() == bb && LD2.lifetimeStart->comesBefore(I);
+
+                auto* LD1_lifetimeStart = LD1.lifetimeStart; // we have to copy LD1.lifetimeStart to avoid clang complaining about LD1 being captured by the lambda
+                bool lifetimeEnd = any_of(LD2.lifetimeEnds, [&](auto* lifetimeEnd) {
+                    return lifetimeEnd->getParent() == bb && LD1_lifetimeStart->comesBefore(lifetimeEnd);
+                });
+
+                if (inflow && outflow)
+                    return true;
+
+                if (inflow && lifetimeEnd)
+                    return true;
+
+                if (outflow && lifetimeStart)
+                    return true;
+
+                if (lifetimeEnd && lifetimeStart)
+                    return true;
+            }
+            else if (I->getParent() == LD2.lifetimeStart->getParent())
             {
                 if (LD2.lifetimeStart->comesBefore(I))
                     return true;

IGC/AdaptorCommon/RayTracing/MergeAllocas.h

@@ -14,6 +14,7 @@ SPDX-License-Identifier: MIT
 
 namespace llvm {
     class BasicBlock;
+    class DominatorTree;
     class Function;
     class Instruction;
     class LoopInfo;
@@ -68,7 +69,14 @@ namespace IGC
             llvm::DenseSet<llvm::BasicBlock*> bbIn;
             llvm::DenseSet<llvm::BasicBlock*> bbOut;
 
-            LivenessData(llvm::Instruction* allocationInstruction, const llvm::SetVector<llvm::Instruction*>& usersOfAllocation, const llvm::LoopInfo& LI, llvm::BasicBlock* userDominatorBlock = nullptr);
+            LivenessData(
+                llvm::Instruction* allocationInstruction,
+                const llvm::SetVector<llvm::Instruction*>& usersOfAllocation,
+                const llvm::LoopInfo& LI,
+                const llvm::DominatorTree& DT,
+                llvm::BasicBlock* userDominatorBlock = nullptr,
+                bool isLifetimeInfinite = false
+            );
 
             bool OverlapsWith(const LivenessData& LD) const;
         };

IGC/Compiler/tests/MergeAllocas/basic-6.ll

@@ -0,0 +1,43 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
+; RUN: igc_opt --opaque-pointers -S --platformbmg --igc-merge-allocas %s | FileCheck %s
+
+; Test interaction between control flow and escaping lifetime
+
+target datalayout = "e-p:64:64:64-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"
+
+declare void @escape(ptr %a)
+
+; Function Attrs: alwaysinline nounwind
+define void @fun() #0 {
+start:
+  %alloca_inf = alloca float, align 8
+  %alloca = alloca float, align 8
+; CHECK: alloca float
+; CHECK-NOT: alloca float
+  br label %altpath
+
+altpath:
+  %a = load float, ptr %alloca
+  br i1 undef, label %altpath1, label %altpath2
+
+altpath1:
+  call void @escape(ptr %alloca_inf)
+  br label %altpathexit
+
+altpath2:
+  br label %altpathexit
+
+altpathexit:
+  ret void
+}
+
+attributes #0 = { nounwind }
+

IGC/Compiler/tests/MergeAllocas/basic-7.ll

@@ -0,0 +1,42 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
+target datalayout = "e-p:64:64:64-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
+
+; Test interaction between control flow and escaping lifetime
+
+declare void @escape(ptr %a)
+
+; Function Attrs: alwaysinline nounwind
+define void @fun() #0 {
+start:
+  %alloca_inf = alloca float
+  %alloca = alloca float
+; CHECK: alloca float
+; CHECK: alloca float
+  br label %altpath
+
+altpath:
+  br i1 undef, label %altpath1, label %altpath2
+
+altpath1:
+  call void @escape(ptr %alloca_inf)
+  br label %altpathexit
+
+altpath2:
+  br label %altpathexit
+
+altpathexit:
+  %a = load float, ptr %alloca
+  ret void
+}
+
+attributes #0 = { nounwind }

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

@@ -1,6 +1,6 @@
 ;=========================== begin_copyright_notice ============================
 ;
-; Copyright (C) 2024 Intel Corporation
+; Copyright (C) 2025 Intel Corporation
 ;
 ; SPDX-License-Identifier: MIT
 ;
@@ -43,14 +43,3 @@ commonret:
 }
 
 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}