reduce: fix loop to selection pass for loops with combined header/continue block (#2480)

* Fix #2478. The fix is to just not try to simplify such loops. 
* Also added `BasicBlock::MergeBlockId()` and `BasicBlock::ContinueBlockId()`. 
* Some minor changes to `structured_loop_to_selection_reduction_opportunity.cpp`. 
* Added test.

source/opt/basic_block.cpp

@@ -184,6 +184,12 @@ uint32_t BasicBlock::MergeBlockIdIfAny() const {
   return mbid;
 }
 
+uint32_t BasicBlock::MergeBlockId() const {
+  uint32_t mbid = MergeBlockIdIfAny();
+  assert(mbid && "Expected block to have a corresponding merge block");
+  return mbid;
+}
+
 uint32_t BasicBlock::ContinueBlockIdIfAny() const {
   auto merge_ii = cend();
   --merge_ii;
@@ -197,6 +203,12 @@ uint32_t BasicBlock::ContinueBlockIdIfAny() const {
   return cbid;
 }
 
+uint32_t BasicBlock::ContinueBlockId() const {
+  uint32_t cbid = ContinueBlockIdIfAny();
+  assert(cbid && "Expected block to have a corresponding continue target");
+  return cbid;
+}
+
 std::ostream& operator<<(std::ostream& str, const BasicBlock& block) {
   str << block.PrettyPrint();
   return str;

source/opt/basic_block.h

@@ -183,10 +183,16 @@ class BasicBlock {
   // block, if any.  If none, returns zero.
   uint32_t MergeBlockIdIfAny() const;
 
+  // Returns MergeBlockIdIfAny() and asserts that it is non-zero.
+  uint32_t MergeBlockId() const;
+
   // Returns the ID of the continue block declared by a merge instruction in
   // this block, if any.  If none, returns zero.
   uint32_t ContinueBlockIdIfAny() const;
 
+  // Returns ContinueBlockIdIfAny() and asserts that it is non-zero.
+  uint32_t ContinueBlockId() const;
+
   // Returns the terminator instruction.  Assumes the terminator exists.
   Instruction* terminator() { return &*tail(); }
   const Instruction* terminator() const { return &*ctail(); }

source/reduce/structured_loop_to_selection_reduction_opportunity.cpp

@@ -23,7 +23,6 @@ namespace reduce {
 
 namespace {
 const uint32_t kMergeNodeIndex = 0;
-const uint32_t kContinueNodeIndex = 1;
 }  // namespace
 
 bool StructuredLoopToSelectionReductionOpportunity::PreconditionHolds() {
@@ -43,15 +42,11 @@ void StructuredLoopToSelectionReductionOpportunity::Apply() {
 
   // (1) Redirect edges that point to the loop's continue target to their
   // closest merge block.
-  RedirectToClosestMergeBlock(
-      loop_construct_header_->GetLoopMergeInst()->GetSingleWordOperand(
-          kContinueNodeIndex));
+  RedirectToClosestMergeBlock(loop_construct_header_->ContinueBlockId());
 
   // (2) Redirect edges that point to the loop's merge block to their closest
   // merge block (which might be that of an enclosing selection, for instance).
-  RedirectToClosestMergeBlock(
-      loop_construct_header_->GetLoopMergeInst()->GetSingleWordOperand(
-          kMergeNodeIndex));
+  RedirectToClosestMergeBlock(loop_construct_header_->MergeBlockId());
 
   // (3) Turn the loop construct header into a selection.
   ChangeLoopToSelection();
@@ -127,12 +122,8 @@ void StructuredLoopToSelectionReductionOpportunity::RedirectEdge(
 
   // original_target_id must either be the merge target or continue construct
   // for the loop being operated on.
-  assert(original_target_id ==
-             loop_construct_header_->GetMergeInst()->GetSingleWordOperand(
-                 kMergeNodeIndex) ||
-         original_target_id ==
-             loop_construct_header_->GetMergeInst()->GetSingleWordOperand(
-                 kContinueNodeIndex));
+  assert(original_target_id == loop_construct_header_->MergeBlockId() ||
+         original_target_id == loop_construct_header_->ContinueBlockId());
 
   auto terminator = context_->cfg()->block(source_id)->terminator();
 
@@ -221,7 +212,7 @@ void StructuredLoopToSelectionReductionOpportunity::ChangeLoopToSelection() {
     const analysis::Bool* bool_type =
         context_->get_type_mgr()->GetRegisteredType(&temp)->AsBool();
     auto const_mgr = context_->get_constant_mgr();
-    auto true_const = const_mgr->GetConstant(bool_type, {true});
+    auto true_const = const_mgr->GetConstant(bool_type, {1});
     auto true_const_result_id =
         const_mgr->GetDefiningInstruction(true_const)->result_id();
     auto original_branch_id = terminator->GetSingleWordOperand(0);

source/reduce/structured_loop_to_selection_reduction_opportunity_finder.cpp

@@ -33,10 +33,9 @@ StructuredLoopToSelectionReductionOpportunityFinder::GetAvailableOpportunities(
   std::set<uint32_t> merge_block_ids;
   for (auto& function : *context->module()) {
     for (auto& block : function) {
-      auto merge_inst = block.GetMergeInst();
-      if (merge_inst) {
-        merge_block_ids.insert(
-            merge_inst->GetSingleWordOperand(kMergeNodeIndex));
+      auto merge_block_id = block.MergeBlockIdIfAny();
+      if (merge_block_id) {
+        merge_block_ids.insert(merge_block_id);
       }
     }
   }
@@ -50,11 +49,19 @@ StructuredLoopToSelectionReductionOpportunityFinder::GetAvailableOpportunities(
         continue;
       }
 
+      uint32_t continue_block_id =
+          loop_merge_inst->GetSingleWordOperand(kContinueNodeIndex);
+
       // Check whether the loop construct's continue target is the merge block
       // of some structured control flow construct.  If it is, we cautiously do
       // not consider applying a transformation.
-      if (merge_block_ids.find(loop_merge_inst->GetSingleWordOperand(
-              kContinueNodeIndex)) != merge_block_ids.end()) {
+      if (merge_block_ids.find(continue_block_id) != merge_block_ids.end()) {
+        continue;
+      }
+
+      // Check whether the loop header block is also the continue target. If it
+      // is, we cautiously do not consider applying a transformation.
+      if (block.id() == continue_block_id) {
         continue;
       }
 

test/reduce/structured_loop_to_selection_test.cpp

@@ -3587,6 +3587,40 @@ TEST(StructuredLoopToSelectionReductionPassTest, LoopyShaderWithOpDecorate) {
   CheckEqual(env, after_op_0, context.get());
 }
 
+TEST(StructuredLoopToSelectionReductionPassTest,
+     LoopWithCombinedHeaderAndContinue) {
+  // A shader containing a loop where the header is also the continue target.
+  // For now, we don't simplify such loops.
+
+  std::string shader = R"(
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %4 "main"
+               OpExecutionMode %4 OriginUpperLeft
+               OpSource ESSL 310
+          %2 = OpTypeVoid
+          %3 = OpTypeFunction %2
+          %6 = OpTypeBool
+         %30 = OpConstantFalse %6
+          %4 = OpFunction %2 None %3
+          %5 = OpLabel
+               OpBranch %10
+         %10 = OpLabel                       ; loop header and continue target
+               OpLoopMerge %12 %10 None
+               OpBranchConditional %30 %10 %12
+         %12 = OpLabel
+               OpReturn
+               OpFunctionEnd
+  )";
+
+  const auto env = SPV_ENV_UNIVERSAL_1_3;
+  const auto context = BuildModule(env, nullptr, shader, kReduceAssembleOption);
+  const auto ops = StructuredLoopToSelectionReductionOpportunityFinder()
+                       .GetAvailableOpportunities(context.get());
+  ASSERT_EQ(0, ops.size());
+}
+
 }  // namespace
 }  // namespace reduce
 }  // namespace spvtools