shader: route nested shared regions

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

@@ -1520,10 +1520,7 @@ bool RouteOneSharedArm(Graph& graph, uint32_t original_block_count, uint32_t out
                        uint32_t route_variable, GotoRouteBlocks& route) {
 	auto* outer = graph.FindBlock(outer_id);
 	if (outer == nullptr || outer->inst_begin == outer->inst_end ||
-	    outer->terminator.kind != TerminatorKind::ConditionalBranch ||
-	    std::ranges::any_of(graph.natural_loops, [outer_id](const auto& loop) {
-		    return Contains(loop.body_blocks, outer_id);
-	    })) {
+	    outer->terminator.kind != TerminatorKind::ConditionalBranch) {
 		return false;
 	}
 
@@ -1553,21 +1550,7 @@ bool RouteOneSharedArm(Graph& graph, uint32_t original_block_count, uint32_t out
 		if (first_arm >= original_block_count || outer_id >= inner_id || inner_id >= first_arm) {
 			continue;
 		}
-		const auto* shared_block = graph.FindBlock(shared);
-		const auto* other_block  = graph.FindBlock(other);
-		if (shared_block == nullptr || other_block == nullptr ||
-		    shared_block->terminator.kind != TerminatorKind::Branch ||
-		    other_block->terminator.kind != TerminatorKind::Branch ||
-		    shared_block->successors.size() != 1u ||
-		    shared_block->successors != other_block->successors) {
-			continue;
-		}
-		const auto in_loop = [&](uint32_t block_id) {
-			return std::ranges::any_of(graph.natural_loops, [&](const auto& loop) {
-				return Contains(loop.body_blocks, block_id);
-			});
-		};
-		if (in_loop(inner_id) || in_loop(shared) || in_loop(other)) {
+		if (graph.FindNearestCommonPostDominator(shared, other) == UINT32_MAX) {
 			continue;
 		}
 
@@ -1592,32 +1575,32 @@ bool RouteOneSharedArm(Graph& graph, uint32_t original_block_count, uint32_t out
 	return false;
 }
 
-bool RouteTerminalSharedSelectionArm(Graph& graph) {
+bool RouteSharedSelectionArms(Graph& graph) {
 	if (graph.irreducible) {
 		return false;
 	}
-	const auto      original_block_count = static_cast<uint32_t>(graph.blocks.size());
-	GotoRouteBlocks route;
-	bool            found = false;
-	for (uint32_t block_id = 0; block_id < original_block_count; block_id++) {
-		GotoRouteBlocks next;
-		if (RouteOneSharedArm(graph, original_block_count, block_id, 0u, next)) {
-			if (found) {
-				// Multiple interacting overlaps remain on the dispatcher path.
-				return false;
+	const auto route_budget = static_cast<uint32_t>(graph.blocks.size());
+	bool       found        = false;
+	for (uint32_t route_variable = 0; route_variable < route_budget; route_variable++) {
+		const auto block_count = static_cast<uint32_t>(graph.blocks.size());
+		bool       routed      = false;
+		for (uint32_t block_id = 0; block_id < block_count; block_id++) {
+			GotoRouteBlocks route;
+			if (!RouteOneSharedArm(graph, block_count, block_id, route_variable, route)) {
+				continue;
 			}
-			found = true;
-			route = next;
+			PlaceGotoRouteBlocks(graph, block_count, route);
 			RebuildPredecessors(graph);
+			RecomputeAnalyses(graph);
+			found  = true;
+			routed = true;
+			break;
+		}
+		if (!routed) {
+			return found;
 		}
 	}
-	if (!found) {
-		return false;
-	}
-	PlaceGotoRouteBlocks(graph, original_block_count, route);
-	RebuildPredecessors(graph);
-	RecomputeAnalyses(graph);
-	return true;
+	return false;
 }
 
 } // namespace
@@ -1939,7 +1922,7 @@ bool Structurize(Graph& graph, std::string* error) {
 	if (StructurizeImpl(graph, error)) {
 		return true;
 	}
-	if (!RouteTerminalSharedSelectionArm(original)) {
+	if (!RouteSharedSelectionArms(original)) {
 		return false;
 	}
 	Graph       failed_graph = std::move(graph);

tests/shaderCfgTests.cpp

@@ -7109,6 +7109,74 @@ void TestNewShaderRecompilerCfgExecSccSharedArm() {
   }
 }
 
+void TestNewShaderRecompilerCfgLoopSharedRegion() {
+  const uint32_t shader[] = {
+      EncodeSopc(0x06, 0, 0),      // loop condition
+      EncodeSopp(0x04, 14),        // loop exit -> end
+      EncodeSopc(0x06, 1, 1),      // outer condition
+      EncodeSopp(0x04, 2),         // outer -> shared region or inner
+      EncodeSopc(0x06, 2, 2),      // inner condition
+      EncodeSopp(0x04, 6),         // inner -> common tail or shared region
+      EncodeSopc(0x06, 3, 3),      // shared region condition
+      EncodeSopp(0x04, 2),         // shared region -> right or left
+      EncodeSMovB32(4, 129),       // shared left work
+      EncodeSopp(0x02, 2),         // shared left -> common tail
+      EncodeSMovB32(5, 129),       // shared right work
+      EncodeSopp(0x02, 0),         // shared right -> common tail
+      EncodeSop2(0x00, 6, 6, 129), // common tail work
+      EncodeSopp(0x02, 0),         // common tail -> continue
+      EncodeSop2(0x00, 7, 7, 129), // continue work
+      EncodeSopp(0x02, 0xfff0u),   // backedge -> loop header
+      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_block_count = graph.blocks.size();
+  const auto original_coverage =
+      CfgInstructionCoverage(graph, decoded.instructions.size());
+  Check(graph.natural_loops.size() == 1u,
+        "loop shared-region fixture has the wrong native CFG");
+  const bool structured = ShaderRecompiler::CFG::Structurize(graph, &error);
+  Check(structured, error.c_str());
+  const auto *loop_header = graph.FindBlockByPc(0x00u);
+  Check(graph.natural_loops.size() == 1u && graph.back_edges.size() == 1u &&
+            loop_header != nullptr && loop_header->terminator.loop_header &&
+            loop_header->terminator.continue_block != UINT32_MAX,
+        "loop shared-region routing did not preserve the natural loop");
+
+  uint32_t route_selects = 0;
+  uint32_t route_sets = 0;
+  for (const auto &block : graph.blocks) {
+    route_selects += block.terminator.condition ==
+                     ShaderRecompiler::CFG::BranchCondition::GotoVariable;
+    route_sets += block.terminator.goto_value >= 0;
+  }
+  Check(route_selects == 1u && route_sets == 3u &&
+            graph.blocks.size() >= original_block_count + 5u &&
+            CfgInstructionCoverage(graph, decoded.instructions.size()) ==
+                original_coverage,
+        "loop shared region was not routed without semantic duplication");
+
+  auto options = MakeCompileOptions(ShaderType::Compute);
+  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, 246) == 1u &&
+            SpirvInstructionOpcodeCount(result.spirv, 251) == 0u,
+        "loop shared region unexpectedly selected dispatcher fallback");
+  CheckSpirvBinaryValidates(result.spirv);
+}
+
 void TestNewShaderRecompilerCfgOverlappingEarlyExitLadder() {
   const uint32_t shader[] = {
       EncodeSopc(0x06, 0, 0), // block 0
@@ -9409,6 +9477,7 @@ int main() {
   TestNewShaderRecompilerCfgDuplicateMergeStructuredSplit();
   TestNewShaderRecompilerCfgNestedEarlyExitLoopForwarders();
   TestNewShaderRecompilerCfgExecSccSharedArm();
+  TestNewShaderRecompilerCfgLoopSharedRegion();
   TestNewShaderRecompilerCfgOverlappingEarlyExitLadder();
   TestNewShaderRecompilerCfgExternallyEnteredSelectionDispatcher();
   TestNewShaderRecompilerCfgIrreducibleDispatcher();