Fix lifetime leaking incorrectly interacting with loops.

A case was found where, for a given alloca,
* all users were contained in a single loop (so there are no cases where the lifetime contains an exiting edge),
* one of the instructions was leaking the lifetime

The result was that the entire loop was included in the lifetime set, but the preheader wasn't. This ended up crashing because there was no single block that could be marked as starting the lifetime.
The fix is to make liveness leak analyzing block work on it's own separate sets of basic block, then merge the result with the actual bbIn and bbOut sets.

Author: jaladreips

IGC/AdaptorCommon/RayTracing/MergeAllocas.cpp

@@ -224,9 +224,8 @@ AllocationBasedLivenessAnalysis::LivenessData* AllocationBasedLivenessAnalysis::
     auto* DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto* LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
 
-    bool hasNoLifetimeEnd = false;
-
     SetVector<Instruction*> allUsers;
+    SetVector<Instruction*> lifetimeLeakingUsers;
     SmallVector<Use*> worklist;
 
     for (auto& use : I->uses())
@@ -266,31 +265,31 @@ AllocationBasedLivenessAnalysis::LivenessData* AllocationBasedLivenessAnalysis::
 
             break;
         case Instruction::PtrToInt:
-            hasNoLifetimeEnd = true;
+            lifetimeLeakingUsers.insert(II);
             break;
         case Instruction::Store:
             {
                 auto* storeI = cast<StoreInst>(II);
                 if (storeI->getValueOperand() == cast<Value>(use))
-                    hasNoLifetimeEnd = true;
+                    lifetimeLeakingUsers.insert(II);
             }
             break;
         case Instruction::Call:
             {
                 auto* callI = cast<CallInst>(II);
                 if (!callI->doesNotCapture(use->getOperandNo()))
-                    hasNoLifetimeEnd = true;
+                    lifetimeLeakingUsers.insert(II);
             }
             break;
         case Instruction::Load:
             break;
         default: // failsafe for handling "unapproved" instructions
-            hasNoLifetimeEnd = true;
+            lifetimeLeakingUsers.insert(II);
             break;
         }
     }
 
-    return new LivenessData(I, allUsers, *LI, *DT, commonDominator, hasNoLifetimeEnd);
+    return new LivenessData(I, std::move(allUsers), *LI, *DT, commonDominator, std::move(lifetimeLeakingUsers));
 }
 
 template<typename range>
@@ -327,8 +326,14 @@ static inline void doWorkLoop(
 
 }
 
-AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocationInstruction, const SetVector<Instruction*>& usersOfAllocation, const LoopInfo& LI, const DominatorTree& DT, BasicBlock* userDominatorBlock, bool isLifetimeInfinite)
-{
+AllocationBasedLivenessAnalysis::LivenessData::LivenessData(
+    Instruction* allocationInstruction,
+    SetVector<Instruction*>&& usersOfAllocation,
+    const LoopInfo& LI,
+    const DominatorTree& DT,
+    BasicBlock* userDominatorBlock,
+    SetVector<Instruction*>&& lifetimeLeakingUsers
+) {
     if (!userDominatorBlock)
         userDominatorBlock = allocationInstruction->getParent();
 
@@ -363,25 +368,32 @@ AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocat
     );
 
     // handle infinite lifetime
-    if (isLifetimeInfinite)
+    if (!lifetimeLeakingUsers.empty())
     {
         // traverse all the successors until there are no left.
-        auto bbInOnly = bbIn;
-        set_subtract(bbInOnly, bbOut);
-
-        for (auto* bb : bbInOnly)
-            worklist.push_back(bb);
+        decltype(bbIn) leakingbbIn;
+        decltype(bbOut) leakingbbOut;
 
-        // in case the only use is the one that causes lifetime escape
-        worklist.push_back(userDominatorBlock);
+        for (auto* I : lifetimeLeakingUsers)
+            worklist.push_back(I->getParent());
 
         doWorkLoop<llvm::succ_range>(
             worklist,
-            bbOut,
-            bbIn,
+            leakingbbOut,
+            leakingbbIn,
             [&](auto* currbb) { return llvm::successors(currbb); },
             [&](auto* currbb) { return false; }
         );
+
+        // add terminators to users, so we can later add them to our lifetimeEnd vector
+        auto leakingbbOnlyIn = leakingbbIn;
+        set_subtract(leakingbbOnlyIn, leakingbbOut);
+
+        for (auto* bb : leakingbbOnlyIn)
+            usersOfAllocation.insert(bb->getTerminator());
+
+        set_union(bbIn, leakingbbIn);
+        set_union(bbOut, leakingbbOut);
     }
 
     // if the lifetime escapes any loop, we should make sure all the loops blocks are included
@@ -469,7 +481,7 @@ AllocationBasedLivenessAnalysis::LivenessData::LivenessData(Instruction* allocat
         {
             for (auto& I : llvm::reverse(*bb))
             {
-                if (usersOfAllocation.contains(&I) || isLifetimeInfinite)
+                if (usersOfAllocation.contains(&I))
                 {
                     lifetimeEnds.push_back(&I);
                     break;

IGC/AdaptorCommon/RayTracing/MergeAllocas.h

@@ -56,11 +56,11 @@ namespace IGC
 
             LivenessData(
                 llvm::Instruction* allocationInstruction,
-                const llvm::SetVector<llvm::Instruction*>& usersOfAllocation,
+                llvm::SetVector<llvm::Instruction*>&& usersOfAllocation,
                 const llvm::LoopInfo& LI,
                 const llvm::DominatorTree& DT,
                 llvm::BasicBlock* userDominatorBlock = nullptr,
-                bool isLifetimeInfinite = false
+                llvm::SetVector<llvm::Instruction*>&& lifetimeLeakingUsers = {}
             );
 
             bool OverlapsWith(const LivenessData& LD) const;

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

@@ -0,0 +1,42 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 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
+
+; Test if lifetime leaking instruction will interact correctly with loops that contain the entirety of the lifetime
+; CHECK-NOT: error
+
+define void @test() #0 {
+entry:
+  %alloca1 = alloca i32
+  br label %looppreheader
+
+looppreheader:
+  br label %loopheader
+
+loopheader:
+  %p2i = ptrtoint ptr %alloca1 to i64
+  br i1 false, label %loopbody, label %if.end
+
+loopbody:
+  %a = load i32, ptr %alloca1
+  store i32 10, ptr %alloca1
+  br i1 false, label %exit, label %if.end
+
+if.end:
+  %b = load i32, ptr %alloca1
+  br label %loopheader
+
+exit:
+  ret void
+}
+
+attributes #0 = { nounwind readnone }