Use structured order to unroll loops. (#3443)

Fixes #3441

source/opt/loop_descriptor.cpp

@@ -485,10 +485,27 @@ void Loop::ComputeLoopStructuredOrder(
 
   if (include_pre_header && GetPreHeaderBlock())
     ordered_loop_blocks->push_back(loop_preheader_);
-  cfg.ForEachBlockInReversePostOrder(
-      loop_header_, [ordered_loop_blocks, this](BasicBlock* bb) {
-        if (IsInsideLoop(bb)) ordered_loop_blocks->push_back(bb);
-      });
+
+  bool is_shader =
+      context_->get_feature_mgr()->HasCapability(SpvCapabilityShader);
+  if (!is_shader) {
+    cfg.ForEachBlockInReversePostOrder(
+        loop_header_, [ordered_loop_blocks, this](BasicBlock* bb) {
+          if (IsInsideLoop(bb)) ordered_loop_blocks->push_back(bb);
+        });
+  } else {
+    // If this is a shader, it is possible that there are unreachable merge and
+    // continue blocks that must be copied to retain the structured order.
+    // The structured order will include these.
+    std::list<BasicBlock*> order;
+    cfg.ComputeStructuredOrder(loop_header_->GetParent(), loop_header_, &order);
+    for (BasicBlock* bb : order) {
+      if (bb == GetMergeBlock()) {
+        break;
+      }
+      ordered_loop_blocks->push_back(bb);
+    }
+  }
   if (include_merge && GetMergeBlock())
     ordered_loop_blocks->push_back(loop_merge_);
 }

test/opt/loop_optimizations/loop_fission.cpp

@@ -2203,11 +2203,11 @@ OpBranchConditional %41 %42 %54
 %45 = OpAccessChain %18 %4 %39
 OpStore %45 %44
 %46 = OpIEqual %12 %39 %19
-OpSelectionMerge %47 None
-OpBranchConditional %46 %51 %47
+OpSelectionMerge %48 None
+OpBranchConditional %46 %47 %48
 %47 = OpLabel
 OpBranch %52
-%51 = OpLabel
+%48 = OpLabel
 OpBranch %52
 %52 = OpLabel
 %53 = OpIAdd %8 %39 %19
@@ -2242,16 +2242,16 @@ OpBranch %21
 OpReturn
 OpFunctionEnd
 )";
-  // clang-format on
-  std::unique_ptr<IRContext> context =
-      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, source,
-                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
-  Module* module = context->module();
-  EXPECT_NE(nullptr, module) << "Assembling failed for shader:\n"
-                             << source << std::endl;
+// clang-format on
+std::unique_ptr<IRContext> context =
+    BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, source,
+                SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+Module* module = context->module();
+EXPECT_NE(nullptr, module) << "Assembling failed for shader:\n"
+                           << source << std::endl;
 
-  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER);
-  SinglePassRunAndCheck<LoopFissionPass>(source, expected, true);
+SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER);
+SinglePassRunAndCheck<LoopFissionPass>(source, expected, true);
 }
 
 /*

test/opt/loop_optimizations/unroll_simple.cpp

@@ -2998,6 +2998,76 @@ TEST_F(PassClassTest, OpPhiSelfReference) {
                                            kUnrollFactor);
 }
 
+// Test that a loop containing an unreachable merge block can still be unrolled
+// correctly.
+TEST_F(PassClassTest, UnreachableMerge) {
+  const std::string text = R"(
+; Identify the first iteration of the unrolled loop, and make sure it contains
+; the unreachable merge block.
+; The first SelectionMerge corresponds to the original loop merge.
+; The second is the branch in the loop.
+; CHECK: OpSelectionMerge {{%\w+}} None
+; CHECK: OpSelectionMerge [[unrch1:%\w+]] None
+; CHECK: [[unrch1]] = OpLabel
+; CHECK-NEXT: OpUnreachable
+; Identify the second iteration of the unrolled loop, and make sure it contains
+; the unreachable merge block.
+; The first SelectionMerge corresponds to the original loop merge
+; The second is the branch in the loop.
+; CHECK: OpSelectionMerge {{%\w+}} None
+; CHECK: OpSelectionMerge [[unrch2:%\w+]] None
+; CHECK: [[unrch2]] = OpLabel
+; CHECK-NEXT: OpUnreachable
+
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main"
+               OpExecutionMode %main LocalSize 64 1 1
+               OpSource HLSL 600
+               OpName %main "main"
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+     %uint_2 = OpConstant %uint 2
+     %uint_1 = OpConstant %uint 1
+       %bool = OpTypeBool
+       %void = OpTypeVoid
+         %18 = OpTypeFunction %void
+       %main = OpFunction %void None %18
+         %23 = OpLabel
+               OpBranch %24
+         %24 = OpLabel
+         %28 = OpPhi %uint %uint_0 %23 %29 %27
+         %30 = OpULessThan %bool %28 %uint_2
+               OpLoopMerge %31 %27 Unroll
+               OpBranchConditional %30 %32 %31
+         %32 = OpLabel
+               OpSelectionMerge %33 None
+               OpSwitch %uint_0 %34
+         %34 = OpLabel
+         %35 = OpUndef %bool
+               OpSelectionMerge %36 None
+               OpBranchConditional %35 %37 %38
+         %38 = OpLabel
+               OpBranch %33
+         %37 = OpLabel
+               OpBranch %33
+         %36 = OpLabel
+               OpUnreachable
+         %33 = OpLabel
+               OpBranch %27
+         %27 = OpLabel
+         %29 = OpIAdd %uint %28 %uint_1
+               OpBranch %24
+         %31 = OpLabel
+               OpReturn
+               OpFunctionEnd
+  )";
+
+  const bool kFullyUnroll = true;
+  const uint32_t kUnrollFactor = 0;
+  SinglePassRunAndMatch<opt::LoopUnroller>(text, true, kFullyUnroll,
+                                           kUnrollFactor);
+}
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools