shader: structurize nested early-exit tails

src/graphics/shader/recompiler/frontend/cfg/ShaderCFG.cpp

@@ -1789,61 +1789,78 @@ bool RouteOneSharedArm(Graph& graph, uint32_t original_block_count, uint32_t out
 			return true;
 		}
 
-		// H0 -> continuation/H1, H1 -> exit/body, body -> continuation. Route both
+		// H0 -> arm/H1, H1 -> exit/body, and both arms reach a common continuation.
-		// destinations after joining H0 and H1 so neither selection has a nonlocal edge.
+		// Route exit/continuation after joining H0 and H1 so neither has a nonlocal edge.
 		for (const bool exit_is_true: {true, false}) {
 			const auto other =
 			    exit_is_true ? inner->terminator.true_block : inner->terminator.false_block;
 			const auto body =
 			    exit_is_true ? inner->terminator.false_block : inner->terminator.true_block;
 			if (other == outer_id || other == inner_id || other == shared || body == outer_id ||
-			    body == inner_id || body == shared || !graph.Dominates(inner_id, body) ||
+			    body == inner_id || body == shared || !graph.Dominates(inner_id, body)) {
-			    !graph.PostDominates(shared, body)) {
 				continue;
 			}
 
-			std::vector<uint32_t> inner_shared_predecessors;
+			const auto continuation = graph.FindNearestCommonPostDominator(shared, body);
-			const auto*           shared_block = graph.FindBlock(shared);
+			const auto* continuation_block = graph.FindBlock(continuation);
-			for (const auto predecessor: shared_block->predecessors) {
+			if (continuation_block == nullptr || continuation == other) {
-				if (predecessor != inner_id && graph.Dominates(inner_id, predecessor)) {
+				continue;
-					inner_shared_predecessors.push_back(predecessor);
+			}
+			std::vector<uint32_t> outer_predecessors;
+			std::vector<uint32_t> inner_predecessors;
+			bool                  external_predecessor = false;
+			for (const auto predecessor: continuation_block->predecessors) {
+				if (predecessor == outer_id || graph.Dominates(shared, predecessor)) {
+					outer_predecessors.push_back(predecessor);
+				} else if (predecessor == inner_id || graph.Dominates(body, predecessor)) {
+					inner_predecessors.push_back(predecessor);
+				} else {
+					external_predecessor = true;
 				}
 			}
-			const auto first_arm = std::min(shared, other);
+			const auto first_arm = std::min(continuation, other);
-			if (inner_shared_predecessors.empty() || first_arm >= original_block_count ||
+			if (outer_predecessors.empty() || inner_predecessors.empty() ||
+			    external_predecessor || first_arm >= original_block_count ||
 			    outer_id >= inner_id || inner_id >= first_arm ||
-			    graph.FindNearestCommonPostDominator(shared, other) == UINT32_MAX) {
+			    graph.FindNearestCommonPostDominator(continuation, other) == UINT32_MAX) {
 				continue;
 			}
 
-			const auto route_select = AppendGotoSelectBlock(graph, route_variable, other, shared);
+			const auto route_select =
+			    AppendGotoSelectBlock(graph, route_variable, other, continuation);
 			const auto inner_merge  = AppendSyntheticBranchBlock(graph, route_select);
-			const auto outer_shared =
+			const auto outer_continue =
 			    AppendGotoSetBlock(graph, route_variable, false, route_select);
-			const auto inner_shared = AppendGotoSetBlock(graph, route_variable, false, inner_merge);
+			const auto inner_continue =
-			const auto inner_other  = AppendGotoSetBlock(graph, route_variable, true, inner_merge);
+			    AppendGotoSetBlock(graph, route_variable, false, inner_merge);
+			const auto inner_other = AppendGotoSetBlock(graph, route_variable, true, inner_merge);
 
-			outer                         = graph.FindBlock(outer_id);
+			auto* mutable_inner = graph.FindBlock(inner_id);
-			auto* mutable_inner           = graph.FindBlock(inner_id);
-			outer->terminator.true_block  = inner_is_true ? inner_id : outer_shared;
-			outer->terminator.false_block = inner_is_true ? outer_shared : inner_id;
-			outer->successors = {outer->terminator.true_block, outer->terminator.false_block};
 			if (exit_is_true) {
 				mutable_inner->terminator.true_block = inner_other;
 			} else {
 				mutable_inner->terminator.false_block = inner_other;
 			}
 			ReplaceValue(mutable_inner->successors, other, inner_other);
-			for (const auto predecessor: inner_shared_predecessors) {
+			for (const auto predecessor: outer_predecessors) {
 				auto* predecessor_block = graph.FindBlock(predecessor);
-				ReplaceValue(predecessor_block->successors, shared, inner_shared);
+				ReplaceValue(predecessor_block->successors, continuation, outer_continue);
-				ReplaceTerminatorTarget(predecessor_block->terminator, shared, inner_shared);
+				ReplaceTerminatorTarget(predecessor_block->terminator, continuation,
+				                        outer_continue);
+			}
+			for (const auto predecessor: inner_predecessors) {
+				auto* predecessor_block = graph.FindBlock(predecessor);
+				ReplaceValue(predecessor_block->successors, continuation, inner_continue);
+				ReplaceTerminatorTarget(predecessor_block->terminator, continuation,
+				                        inner_continue);
 			}
 
-			std::vector<uint32_t> route_blocks = {inner_shared, inner_other, inner_merge,
+			std::vector<uint32_t> route_blocks = {inner_continue, inner_other, inner_merge,
-			                                      outer_shared};
+			                                      outer_continue};
-			AppendEnclosingRouteBlocks(graph, original_block_count, route_variable, shared, other,
+			if (continuation == shared) {
-			                           route_select, outer_id, route_blocks);
+				AppendEnclosingRouteBlocks(graph, original_block_count, route_variable, shared,
+				                           other, route_select, outer_id, route_blocks);
+			}
 			route_blocks.push_back(route_select);
 			route = {first_arm, std::move(route_blocks)};
 			return true;

tests/shaderCfgTests.cpp

@@ -7684,6 +7684,54 @@ void TestNewShaderRecompilerCfgExecSccSharedArm() {
   }
 }
 
+void TestNewShaderRecompilerCfgNestedTailEarlyExit() {
+  const uint32_t shader[] = {
+      EncodeSopc(0x06, 0, 0), // outer condition
+      EncodeSopp(0x04, 4),    // outer -> right arm or inner condition
+      EncodeSopc(0x06, 1, 1), // inner early-exit condition
+      EncodeSopp(0x04, 4),    // inner -> exit or left arm
+      EncodeSMovB32(2, 129),  // left arm
+      EncodeSopp(0x02, 1),    // left arm -> common tail
+      EncodeSMovB32(3, 129),  // right arm
+      EncodeSMovB32(4, 129),  // common tail
+      0xbf810000u,
+  };
+
+  ShaderRecompiler::Decoder::Program decoded;
+  std::string error;
+  Check(ShaderRecompiler::Decoder::DecodeProgram(std::span{shader}, decoded,
+                                                 &error),
+        error.c_str());
+  ShaderRecompiler::CFG::Graph graph;
+  Check(ShaderRecompiler::CFG::BuildGraph(decoded, graph, &error),
+        error.c_str());
+  const auto original_coverage =
+      CfgInstructionCoverage(graph, decoded.instructions.size());
+  Check(ShaderRecompiler::CFG::Structurize(graph, &error), error.c_str());
+  Check(CfgInstructionCoverage(graph, decoded.instructions.size()) ==
+            original_coverage,
+        "nested-tail routing changed semantic instruction coverage");
+  Check(std::ranges::count_if(graph.blocks, [](const auto &block) {
+          return block.terminator.condition ==
+                 ShaderRecompiler::CFG::BranchCondition::GotoVariable;
+        }) == 1u &&
+            std::ranges::count_if(graph.blocks, [](const auto &block) {
+              return block.terminator.goto_value >= 0;
+            }) == 3u,
+        "nested-tail early exit did not use typed route state");
+
+  auto options = MakeCompileOptions(ShaderType::Pixel);
+  options.dump_ir = true;
+  ShaderRecompiler::CompileResult result;
+  Check(ShaderRecompiler::TryRecompile(shader, options, result, &error),
+        error.c_str());
+  Check(!result.program.dispatcher_fallback &&
+            Common::ContainsStr(result.ir_dump, "mode=structured") &&
+            SpirvInstructionOpcodeCount(result.spirv, 251) == 0u,
+        "nested-tail early exit selected dispatcher control flow");
+  CheckSpirvBinaryValidates(result.spirv);
+}
+
 void TestNewShaderRecompilerCfgLoopSharedRegion() {
   const uint32_t shader[] = {
       EncodeSopc(0x06, 0, 0),      // loop condition
@@ -11582,6 +11630,7 @@ int main() {
   TestNewShaderRecompilerCfgDuplicateMergeStructuredSplit();
   TestNewShaderRecompilerCfgNestedEarlyExitLoopForwarders();
   TestNewShaderRecompilerCfgExecSccSharedArm();
+  TestNewShaderRecompilerCfgNestedTailEarlyExit();
   TestNewShaderRecompilerCfgLoopSharedRegion();
   TestNewShaderRecompilerCfgOverlappingEarlyExitLadder();
   TestNewShaderRecompilerCfgNestedEarlyExitSharedTerminal();