From 46e4b5a3c8b6293732fd5b071a600c13cc4adb56 Mon Sep 17 00:00:00 2001
From: Hans-Kristian Arntzen <post@arntzen-software.no>
Date: Tue, 7 Jun 2022 14:58:48 +0200
Subject: [PATCH] Fix control flow bug where we missed continue;

Case which caused failure:

if (cond)
{
    continue;
}
break;

Only allow tracing from inner selections if the outer header never
merges execution.
---
 ...loop-break-merge-after-inner-continue.comp | 22 +++++++++++++
 ...loop-break-merge-after-inner-continue.comp | 21 ++++++++++++
 spirv_cfg.cpp                                 | 33 +++++++++++++++----
 3 files changed, 70 insertions(+), 6 deletions(-)
 create mode 100644 reference/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp
 create mode 100644 shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp

diff --git a/reference/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp b/reference/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp
new file mode 100644
index 00000000..5b4cb886
--- /dev/null
+++ b/reference/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp
@@ -0,0 +1,22 @@
+#version 450
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0, std430) buffer STO
+{
+    uint data[];
+} ssbo;
+
+void main()
+{
+    while (true)
+    {
+        ssbo.data[0]++;
+        if (ssbo.data[2] != 0u)
+        {
+            ssbo.data[5]++;
+            continue;
+        }
+        break;
+    }
+}
+
diff --git a/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp b/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp
new file mode 100644
index 00000000..e916ab24
--- /dev/null
+++ b/shaders-no-opt/comp/loop-break-merge-after-inner-continue.comp
@@ -0,0 +1,21 @@
+#version 450
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) buffer STO
+{
+  uint data[];
+} ssbo;
+
+void main()
+{
+  while(true)
+  {
+    ssbo.data[0] += 1;
+    if (bool(ssbo.data[2]))
+    {
+      ssbo.data[5] += 1;
+      continue;
+    }
+    break;
+  }
+}
diff --git a/spirv_cfg.cpp b/spirv_cfg.cpp
index a938634e..93299479 100644
--- a/spirv_cfg.cpp
+++ b/spirv_cfg.cpp
@@ -306,15 +306,36 @@ bool CFG::node_terminates_control_flow_in_sub_graph(BlockID from, BlockID to) co
 
 		bool true_path_ignore = false;
 		bool false_path_ignore = false;
-		if (ignore_block_id && dom.terminator == SPIRBlock::Select)
+
+		bool merges_to_nothing = dom.merge == SPIRBlock::MergeNone ||
+		                         (dom.merge == SPIRBlock::MergeSelection && dom.next_block &&
+		                          compiler.get<SPIRBlock>(dom.next_block).terminator == SPIRBlock::Unreachable) ||
+		                         (dom.merge == SPIRBlock::MergeLoop && dom.merge_block &&
+		                          compiler.get<SPIRBlock>(dom.merge_block).terminator == SPIRBlock::Unreachable);
+
+		if (dom.self == from || merges_to_nothing)
 		{
-			auto &true_block = compiler.get<SPIRBlock>(dom.true_block);
-			auto &false_block = compiler.get<SPIRBlock>(dom.false_block);
-			auto &ignore_block = compiler.get<SPIRBlock>(ignore_block_id);
-			true_path_ignore = compiler.execution_is_branchless(true_block, ignore_block);
-			false_path_ignore = compiler.execution_is_branchless(false_block, ignore_block);
+			// We can only ignore inner branchy paths if there is no merge,
+			// i.e. no code is generated afterwards. E.g. this allows us to elide continue:
+			// for (;;) { if (cond) { continue; } else { break; } }.
+			// Codegen here in SPIR-V will be something like either no merge if one path directly breaks, or
+			// we merge to Unreachable.
+			if (ignore_block_id && dom.terminator == SPIRBlock::Select)
+			{
+				auto &true_block = compiler.get<SPIRBlock>(dom.true_block);
+				auto &false_block = compiler.get<SPIRBlock>(dom.false_block);
+				auto &ignore_block = compiler.get<SPIRBlock>(ignore_block_id);
+				true_path_ignore = compiler.execution_is_branchless(true_block, ignore_block);
+				false_path_ignore = compiler.execution_is_branchless(false_block, ignore_block);
+			}
 		}
 
+		// Cases where we allow traversal. This serves as a proxy for post-dominance in a loop body.
+		// TODO: Might want to do full post-dominance analysis, but it's a lot of churn for something like this ...
+		// - We're the merge block of a selection construct. Jump to header.
+		// - We're the merge block of a loop. Jump to header.
+		// - Direct branch. Trivial.
+		// - Allow cases inside a branch if the header cannot merge execution before loop exit.
 		if ((dom.merge == SPIRBlock::MergeSelection && dom.next_block == to) ||
 		    (dom.merge == SPIRBlock::MergeLoop && dom.merge_block == to) ||
 		    (dom.terminator == SPIRBlock::Direct && dom.next_block == to) ||