Update the continue target in merge return. (#2249)

When we are predicating the continue target for a loop, it can no longer
be the continue target because it will have a branch that exits the loop
and is not the bach edge.  The continue target will have to be the
target of that branch that is still in the loop.

Fixes #2211.

source/opt/merge_return_pass.cpp

@@ -342,24 +342,27 @@ bool MergeReturnPass::PredicateBlocks(
   while (block != nullptr && block != final_return_block_) {
     if (!predicated->insert(block).second) break;
     // Skip structured subgraphs.
-    BasicBlock* next = nullptr;
     assert(state->InLoop() && "Should be in the dummy loop at the very least.");
-    next = context()->get_instr_block(state->LoopMergeId());
-    while (state->LoopMergeId() == next->id()) {
+    Instruction* current_loop_merge_inst = state->LoopMergeInst();
+    uint32_t merge_block_id =
+        current_loop_merge_inst->GetSingleWordInOperand(0);
+    while (state->LoopMergeId() == merge_block_id) {
       state++;
     }
-    if (!BreakFromConstruct(block, next, predicated, order)) {
+    if (!BreakFromConstruct(block, predicated, order,
+                            current_loop_merge_inst)) {
       return false;
     }
-    block = next;
+    block = context()->get_instr_block(merge_block_id);
   }
   return true;
 }
 
 bool MergeReturnPass::BreakFromConstruct(
-    BasicBlock* block, BasicBlock* merge_block,
-    std::unordered_set<BasicBlock*>* predicated,
-    std::list<BasicBlock*>* order) {
+    BasicBlock* block, std::unordered_set<BasicBlock*>* predicated,
+    std::list<BasicBlock*>* order, Instruction* loop_merge_inst) {
+  assert(loop_merge_inst->opcode() == SpvOpLoopMerge &&
+         "loop_merge_inst must be a loop merge instruction.");
   // Make sure the CFG is build here.  If we don't then it becomes very hard
   // to know which new blocks need to be updated.
   context()->BuildInvalidAnalyses(IRContext::kAnalysisCFG);
@@ -380,6 +383,9 @@ bool MergeReturnPass::BreakFromConstruct(
       return false;
     }
   }
+
+  uint32_t merge_block_id = loop_merge_inst->GetSingleWordInOperand(0);
+  BasicBlock* merge_block = context()->get_instr_block(merge_block_id);
   if (merge_block->GetLoopMergeInst()) {
     cfg()->SplitLoopHeader(merge_block);
   }
@@ -396,6 +402,13 @@ bool MergeReturnPass::BreakFromConstruct(
   BasicBlock* old_body = block->SplitBasicBlock(context(), TakeNextId(), iter);
   predicated->insert(old_body);
 
+  // If |block| was a continue target for a loop |old_body| is now the correct
+  // continue target.
+  if (loop_merge_inst->GetSingleWordInOperand(1) == block->id()) {
+    loop_merge_inst->SetInOperand(1, {old_body->id()});
+    context()->UpdateDefUse(loop_merge_inst);
+  }
+
   // Update |order| so old_block will be traversed.
   InsertAfterElement(block, old_body, order);
 

source/opt/merge_return_pass.h

@@ -219,16 +219,18 @@ class MergeReturnPass : public MemPass {
                        std::list<BasicBlock*>* order);
 
   // Add a conditional branch at the start of |block| that either jumps to
-  // |merge_block| or the original code in |block| depending on the value in
-  // |return_flag_|.
+  // the merge block of |loop_merge_inst| or the original code in |block|
+  // depending on the value in |return_flag_|.  The continue target in
+  // |loop_merge_inst| will be updated if needed.
   //
   // If new blocks that are created will be added to |order|.  This way a call
   // can traverse these new block in structured order.
   //
   // Returns true if successful.
-  bool BreakFromConstruct(BasicBlock* block, BasicBlock* merge_block,
+  bool BreakFromConstruct(BasicBlock* block,
                           std::unordered_set<BasicBlock*>* predicated,
-                          std::list<BasicBlock*>* order);
+                          std::list<BasicBlock*>* order,
+                          Instruction* loop_merge_inst);
 
   // Add an |OpReturn| or |OpReturnValue| to the end of |block|.  If an
   // |OpReturnValue| is needed, the return value is loaded from |return_value_|.

test/opt/pass_merge_return_test.cpp

@@ -1317,6 +1317,57 @@ TEST_F(MergeReturnPassTest, GeneratePhiInOuterLoop) {
   SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
   SinglePassRunAndMatch<MergeReturnPass>(before, false);
 }
+
+TEST_F(MergeReturnPassTest, InnerLoopMergeIsOuterLoopContinue) {
+  const std::string before =
+      R"(
+      ; CHECK: OpLoopMerge
+      ; CHECK-NEXT: OpBranch [[bb1:%\w+]]
+      ; CHECK: [[bb1]] = OpLabel
+      ; CHECK-NEXT: OpBranch [[outer_loop_header:%\w+]]
+      ; CHECK: [[outer_loop_header]] = OpLabel
+      ; CHECK-NEXT: OpLoopMerge [[outer_loop_merge:%\w+]] [[outer_loop_continue:%\w+]] None
+      ; CHECK: [[outer_loop_continue]] = OpLabel
+      ; CHECK-NEXT: OpBranch [[outer_loop_header]]
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %2 "main"
+               OpExecutionMode %2 OriginUpperLeft
+               OpSource ESSL 310
+       %void = OpTypeVoid
+          %4 = OpTypeFunction %void
+       %bool = OpTypeBool
+          %6 = OpTypeFunction %bool
+       %true = OpConstantTrue %bool
+          %2 = OpFunction %void None %4
+          %8 = OpLabel
+          %9 = OpFunctionCall %bool %10
+               OpReturn
+               OpFunctionEnd
+         %10 = OpFunction %bool None %6
+         %11 = OpLabel
+               OpBranch %12
+         %12 = OpLabel
+               OpLoopMerge %13 %14 None
+               OpBranchConditional %true %15 %13
+         %15 = OpLabel
+               OpLoopMerge %14 %16 None
+               OpBranchConditional %true %17 %14
+         %17 = OpLabel
+               OpReturnValue %true
+         %16 = OpLabel
+               OpBranch %15
+         %14 = OpLabel
+               OpBranch %12
+         %13 = OpLabel
+               OpReturnValue %true
+               OpFunctionEnd
+)";
+
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SinglePassRunAndMatch<MergeReturnPass>(before, false);
+}
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools