Fixes infinite loop in ADCE

* Addresses how breaks are indentified to prevent infinite loops when
back to back loop share a merge and header
* Added test to catch the bug

diff --git a/source/opt/aggressive_dead_code_elim_pass.cpp b/source/opt/aggressive_dead_code_elim_pass.cpp
index 8125d57..c384c6f 100644 (file)
--- a/source/opt/aggressive_dead_code_elim_pass.cpp
+++ b/source/opt/aggressive_dead_code_elim_pass.cpp
@@ -190,10 +190,14 @@ void AggressiveDCEPass::ComputeBlock2HeaderMaps(
     std::list<ir::BasicBlock*>& structuredOrder) {
   block2headerBranch_.clear();
   branch2merge_.clear();
+  structured_order_index_.clear();
   std::stack<ir::Instruction*> currentHeaderBranch;
   currentHeaderBranch.push(nullptr);
   uint32_t currentMergeBlockId = 0;
-  for (auto bi = structuredOrder.begin(); bi != structuredOrder.end(); ++bi) {
+  uint32_t index = 0;
+  for (auto bi = structuredOrder.begin(); bi != structuredOrder.end();
+       ++bi, ++index) {
+    structured_order_index_[*bi] = index;
     // If this block is the merge block of the current control construct,
     // we are leaving the current construct so we must update state
     if ((*bi)->id() == currentMergeBlockId) {
@@ -235,26 +239,24 @@ void AggressiveDCEPass::AddBranch(uint32_t labelId, ir::BasicBlock* bp) {
 
 void AggressiveDCEPass::AddBreaksAndContinuesToWorklist(
     ir::Instruction* loopMerge) {
+  ir::BasicBlock* header = context()->get_instr_block(loopMerge);
+  uint32_t headerIndex = structured_order_index_[header];
   const uint32_t mergeId =
       loopMerge->GetSingleWordInOperand(kLoopMergeMergeBlockIdInIdx);
+  ir::BasicBlock* merge = context()->get_instr_block(mergeId);
+  uint32_t mergeIndex = structured_order_index_[merge];
   get_def_use_mgr()->ForEachUser(
-      mergeId, [&loopMerge, this](ir::Instruction* user) {
-        // A branch to the merge block can only be a break if it is nested in
-        // the current loop
+      mergeId, [headerIndex, mergeIndex, this](ir::Instruction* user) {
         if (!user->IsBranch()) return;
-        ir::Instruction* branchInst = user;
-        while (true) {
-          ir::BasicBlock* blk = context()->get_instr_block(branchInst);
-          ir::Instruction* hdrBranch = block2headerBranch_[blk];
-          if (hdrBranch == nullptr) return;
-          ir::Instruction* hdrMerge = branch2merge_[hdrBranch];
-          if (hdrMerge == loopMerge) break;
-          branchInst = hdrBranch;
+        ir::BasicBlock* block = context()->get_instr_block(user);
+        uint32_t index = structured_order_index_[block];
+        if (headerIndex < index && index < mergeIndex) {
+          // This is a break from the loop.
+          AddToWorklist(user);
+          // Add branch's merge if there is one.
+          ir::Instruction* userMerge = branch2merge_[user];
+          if (userMerge != nullptr) AddToWorklist(userMerge);
         }
-        AddToWorklist(user);
-        // Add branch's merge if there is one
-        ir::Instruction* userMerge = branch2merge_[user];
-        if (userMerge != nullptr) AddToWorklist(userMerge);
       });
   const uint32_t contId =
       loopMerge->GetSingleWordInOperand(kLoopMergeContinueBlockIdInIdx);

diff --git a/source/opt/aggressive_dead_code_elim_pass.h b/source/opt/aggressive_dead_code_elim_pass.h
index e3ae3bb..86b784a 100644 (file)
--- a/source/opt/aggressive_dead_code_elim_pass.h
+++ b/source/opt/aggressive_dead_code_elim_pass.h
@@ -158,6 +158,9 @@ class AggressiveDCEPass : public MemPass {
   // of an enclosing construct's header, if one exists.
   std::unordered_map<ir::BasicBlock*, ir::Instruction*> block2headerBranch_;
 
+  // Maps basic block to their index in the structured order traversal.
+  std::unordered_map<ir::BasicBlock*, uint32_t> structured_order_index_;
+
   // Map from branch to its associated merge instruction, if any
   std::unordered_map<ir::Instruction*, ir::Instruction*> branch2merge_;
 

diff --git a/test/opt/aggressive_dead_code_elim_test.cpp b/test/opt/aggressive_dead_code_elim_test.cpp
index b79eed3..61f0593 100644 (file)
--- a/test/opt/aggressive_dead_code_elim_test.cpp
+++ b/test/opt/aggressive_dead_code_elim_test.cpp
@@ -5178,6 +5178,51 @@ OpFunctionEnd
 }
 #endif  // SPIRV_EFFCEE
 
+// Test for #1214
+TEST_F(AggressiveDCETest, LoopHeaderIsAlsoAnotherLoopMerge) {
+  const std::string text = R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %1 "func" %2
+OpExecutionMode %1 OriginUpperLeft
+%void = OpTypeVoid
+%bool = OpTypeBool
+%true = OpConstantTrue %bool
+%uint = OpTypeInt 32 0
+%_ptr_Output_uint = OpTypePointer Output %uint
+%2 = OpVariable %_ptr_Output_uint Output
+%uint_0 = OpConstant %uint 0
+%9 = OpTypeFunction %void
+%1 = OpFunction %void None %9
+%10 = OpLabel
+OpBranch %11
+%11 = OpLabel
+OpLoopMerge %12 %13 None
+OpBranchConditional %true %14 %13
+%14 = OpLabel
+OpStore %2 %uint_0
+OpLoopMerge %15 %16 None
+OpBranchConditional %true %15 %16
+%16 = OpLabel
+OpBranch %14
+%15 = OpLabel
+OpBranchConditional %true %12 %13
+%13 = OpLabel
+OpBranch %11
+%12 = OpLabel
+%17 = OpPhi %uint %uint_0 %15 %uint_0 %18
+OpStore %2 %17
+OpLoopMerge %19 %18 None
+OpBranchConditional %true %19 %18
+%18 = OpLabel
+OpBranch %12
+%19 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndCheck<opt::AggressiveDCEPass>(text, text, true, true);
+}
+
 TEST_F(AggressiveDCETest, BreaksDontVisitPhis) {
   const std::string text = R"(
 OpCapability Shader