graphics: skip compute dispatches with BVH instructions

Co-authored-by: Mac <me@itsmemac.com>

src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp

@@ -204,6 +204,7 @@ struct PipelineCache::ProgramCache {
 		ShaderRecompiler::IR::ResourceSnapshot       resources;
 		ShaderRecompiler::IR::ResourceSpecialization specialization;
 		std::vector<Permutation>                    permutations;
+		bool                                        skip_dispatch = false;
 	};
 
 	struct ProgramKeyHash {
@@ -283,6 +284,9 @@ struct PipelineCache::ProgramCache {
 		lookup_key.code_size       = static_cast<uint32_t>(params.code.size());
 		BuildStageStaticKey(input_info, lookup_key.static_state);
 		auto                                         entry = programs.find(lookup_key);
+		if (entry != programs.end() && entry->second.skip_dispatch) {
+			return {};
+		}
 		const ShaderRecompiler::IR::SrtRuntime       runtime {
 		    .user_data                  = user_data,
 		    .shader_base                = params.Base(),
@@ -348,6 +352,11 @@ struct PipelineCache::ProgramCache {
 			options.wave_size = input_info.wave_size;
 		}
 		auto translated = ShaderRecompiler::TranslateProgram(params.code, options);
+		if (translated.skip_dispatch) {
+			entry = programs.try_emplace(lookup_key, ShaderRecompiler::IR::ResourcePlan {}).first;
+			entry->second.skip_dispatch = true;
+			return {};
+		}
 		if (entry == programs.end()) {
 			entry = programs.try_emplace(lookup_key,
 			    ShaderRecompiler::IR::ExtractResourcePlan(translated.program)).first;

src/graphics/host_gpu/renderer/renderCompute.cpp

@@ -256,6 +256,10 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer,
 	input_info.dispatch_thread_dimensions = use_thread_dimensions;
 	const auto compute_program =
 	    m_context.GetPipelineCache().GetComputeProgram(cs_regs, sh_regs, input_info);
+	if (!compute_program) {
+		// Temporary until RT is implemented.
+		return;
+	}
 	if (use_thread_dimensions) {
 		input_info.dispatch_threads_num[0]    = thread_group_x;
 		input_info.dispatch_threads_num[1]    = thread_group_y;
@@ -413,6 +417,10 @@ void RenderExecutor::DispatchIndirect(uint64_t submit_id, CommandBuffer& buffer,
 	ShaderComputeInputInfo input_info {};
 	const auto compute_program = m_context.GetPipelineCache().GetComputeProgram(
 	    cs_regs, buffer.GetRegisters().GetShaderRegisters(), input_info);
+	if (!compute_program) {
+		// Temporary until RT is implemented.
+		return;
+	}
 	buffer.EndRendering();
 	auto& pipeline = m_context.GetPipelineCache().GetComputePipeline(input_info, compute_program);
 	auto& bindings = m_compute_bindings;

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -20,6 +20,7 @@
 
 #include <algorithm>
 #include <array>
+#include <atomic>
 #include <chrono>
 #include <fmt/format.h>
 #include <map>
@@ -523,6 +524,20 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
 	     GetDumpLabel(options), StageName(options.stage), options.shader_hash,
 	     static_cast<uint64_t>(decoded.instructions.size()), phase_ms());
 
+	// Temporary workaround for games that compile ray-tracing shaders before
+	// the player can select a mode without ray tracing.
+	if (options.stage == ShaderType::Compute && decoded.has_bvh) {
+		static std::atomic_flag warned = ATOMIC_FLAG_INIT;
+		if (!warned.test_and_set(std::memory_order_relaxed)) {
+			const auto& bvh = decoded.instructions.back();
+			Log::WriteToConsoleAndLog(fmt::format(
+			    "Warning: ray tracing is not implemented; skipping compute dispatches containing "
+			    "BVH intersection instructions (shader=0x{:016x}, pc=0x{:08x}, opcode=0x{:02x}).\n",
+			    options.shader_hash, bvh.pc, bvh.opcode_id));
+		}
+		return {.skip_dispatch = true};
+	}
+
 	std::string decoded_dump;
 	if (options.dump_ir) {
 		decoded_dump = Decoder::ProgramToString(decoded);
@@ -615,6 +630,7 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
 CompileResult CompileProgram(TranslateResult translated, const CompileOptions& options,
                              const IR::ResourceSpecialization& specialization,
                              uint32_t push_data_start_dword) {
+	EXIT_IF(translated.skip_dispatch);
 	const auto emit_begin = std::chrono::steady_clock::now();
 	auto& ir = translated.program;
 	IR::ApplyResourceSpecialization(ir, specialization);

src/graphics/shader/recompiler/ShaderRecompiler.h

@@ -28,6 +28,7 @@ struct TranslateResult {
 	IR::Program program;
 	std::string decoded_dump;
 	std::string cfg_dump;
+	bool        skip_dispatch = false;
 };
 
 struct CompileResult {

src/graphics/shader/recompiler/frontend/decode/ShaderDecoder.cpp

@@ -403,6 +403,7 @@ void DecodeProgram(std::span<const uint32_t> code, Program& program) {
 	program.instructions.clear();
 	program.instructions.reserve(code.size());
 	program.code = code;
+	program.has_bvh = false;
 
 	std::vector<bool> branch_targets;
 	for (uint32_t word_index = 0; word_index < code.size();) {
@@ -411,6 +412,10 @@ void DecodeProgram(std::span<const uint32_t> code, Program& program) {
 
 		const auto& inst = program.instructions.back();
 		word_index += inst.word_count;
+		if (inst.family == Family::MIMG && (inst.opcode_id == 0xe6u || inst.opcode_id == 0xe7u)) {
+			program.has_bvh = true;
+			return;
+		}
 
 		if (IsDirectBranch(inst.opcode)) {
 			const auto target_index = inst.branch_target / sizeof(uint32_t);

src/graphics/shader/recompiler/frontend/decode/ShaderDecoder.h

@@ -731,6 +731,7 @@ struct Instruction {
 struct Program {
 	std::span<const uint32_t> code;
 	std::vector<Instruction>  instructions;
+	bool                     has_bvh = false;
 };
 
 // Code spans are trusted to contain complete instructions, valid branch targets, and 32-bit PCs.

tests/ShaderRayTracingTests.inc

@@ -0,0 +1,87 @@
+void TestRayTracingDispatchDetection() {
+  namespace D = ShaderRecompiler::Decoder;
+  const auto check_decoded = [](std::span<const uint32_t> code, bool expected,
+                                const char *message) {
+    D::Program decoded;
+    D::DecodeProgram(code, decoded);
+    Check(decoded.has_bvh == expected, message);
+    if (expected) {
+      const auto &last = decoded.instructions.back();
+      Check(last.family == D::Family::MIMG &&
+                (last.opcode_id == 0xe6u || last.opcode_id == 0xe7u),
+            "decoder continued past the first BVH instruction");
+      const auto result = ShaderRecompiler::TranslateProgram(
+          code, MakeCompileOptions(ShaderType::Compute));
+      Check(result.skip_dispatch && result.program.blocks.empty() && result.cfg_dump.empty(),
+            "BVH compute shader reached IR translation");
+    }
+  };
+  const auto check = [&](std::initializer_list<uint32_t> code, bool expected,
+                         const char *message) {
+    check_decoded(std::span<const uint32_t>{code.begin(), code.size()}, expected, message);
+  };
+
+  check({0xf1989f01u, 0x00081014u, EncodeSopp(0x01)}, true,
+        "captured contiguous BVH instruction did not skip dispatch");
+  check({0xf1989f07u, 0x00071001u, 0x3d3e3c15u, 0x38095851u, 0x00003a39u,
+         EncodeSopp(0x01)}, true,
+        "captured NSA BVH instruction did not skip dispatch");
+  check({0xf19c9f01u, 0x00081014u, EncodeSopp(0x01)}, true,
+        "RDNA2 BVH64 instruction did not skip dispatch");
+  check({0xbe8e210eu, 0xf1989f01u, 0x00081014u, EncodeSopp(0x01)}, true,
+        "unsupported SWAPPC before BVH prevented dispatch detection");
+  check({0xf1989f01u, 0x00081014u, 0xc0000000u, EncodeSopp(0x01)}, true,
+        "unknown instruction after BVH prevented dispatch skipping");
+
+  for (uint32_t opcode : {0xe6u, 0xe7u}) {
+    for (uint32_t nsa = 0; nsa < 4; ++nsa) {
+      for (uint32_t slc = 0; slc < 2; ++slc) {
+        std::array<uint32_t, 6> code{
+            EncodeMimg0(opcode, 0xf, true) | (nsa << 1u) | (slc << 25u),
+            EncodeMimg1(0, 0, 0, 0)};
+        code[2 + nsa] = EncodeSopp(0x01);
+        check_decoded(std::span(code).first(3 + nsa), true,
+                      "BVH cache flags or NSA length prevented dispatch detection");
+      }
+    }
+  }
+  for (uint32_t opcode : {0x66u, 0x67u, 0xe4u, 0xe5u, 0xe8u, 0xe9u}) {
+    check({EncodeMimg0(opcode, 0xf), EncodeMimg1(0, 0, 0, 0), EncodeSopp(0x01)},
+          false, "non-BVH image opcode was incorrectly skipped");
+  }
+
+  check({EncodeSMovB32(0, 255), 0xf1989f01u, EncodeSopp(0x01)}, false,
+        "scalar literal was mistaken for a BVH instruction");
+  check({EncodeMimg0(0x20, 0xf) | 6u, EncodeMimg1(0, 0, 0, 0),
+         0xf1989f01u, 0xf19c9f01u, 0u, EncodeSopp(0x01)}, false,
+        "NSA address words were mistaken for BVH instructions");
+  check({EncodeMimg0(0x20, 0xf), 0xf1989f01u, EncodeSopp(0x01)}, false,
+        "MIMG operand word was mistaken for a BVH instruction");
+  check({EncodeSopp(0x01), 0xf1989f01u, 0x00081014u}, false,
+        "data after the shader end was mistaken for a BVH instruction");
+  check({EncodeMimg0(0x20, 0xf), EncodeMimg1(0, 0, 0, 0), EncodeSopp(0x01)},
+        false, "ordinary image shader was incorrectly skipped");
+
+  // An end instruction on one branch does not end another branch's code.
+  check({EncodeSopp(0x04, 1), EncodeSopp(0x01), 0xf1989f01u, 0x00081014u,
+         EncodeSopp(0x01)}, true,
+        "BVH at a branch target after an early end was not detected");
+
+  const std::array bvh_code{0xf1989f01u, 0x00081014u, 0xf19c9f01u, 0x00081014u,
+                            EncodeSopp(0x01)};
+  const std::array ordinary_code{EncodeSMovB32(0, 255), 0xf1989f01u, EncodeSopp(0x01)};
+  D::Program reused;
+  D::DecodeProgram(bvh_code, reused);
+  Check(reused.has_bvh && reused.instructions.size() == 1,
+        "decoder did not stop at the first BVH instruction");
+  D::DecodeProgram(ordinary_code, reused);
+  Check(!reused.has_bvh, "decoder retained stale BVH metadata");
+  for (const auto stage : {ShaderType::Compute, ShaderType::Vertex, ShaderType::Pixel}) {
+    const auto options = MakeCompileOptions(stage);
+    Check(!ShaderRecompiler::TranslateProgram(ordinary_code, options).skip_dispatch,
+          "BVH-shaped literal marked an ordinary shader for skipping");
+    const auto compiled = RecompileForTest(ordinary_code, options);
+    Check(!compiled.spirv.empty(), "ordinary shader did not compile");
+    CheckSpirvBinaryValidates(compiled.spirv);
+  }
+}

tests/shaderCfgTests.cpp

@@ -13517,6 +13517,8 @@ void TestNewShaderRecompilerSpirvSizeBaselines() {
   CheckSpirvPhiParents(dispatcher_result.spirv);
 }
 
+#include "ShaderRayTracingTests.inc"
+
 } // namespace
 } // namespace Libs::Graphics
 
@@ -13524,6 +13526,7 @@ int main() {
   using namespace Libs::Graphics;
 
   EnsureConfigInitialized();
+  TestRayTracingDispatchDetection();
   TestResourceDescriptorClassification();
   TestShaderBufferResourceSize();
   TestNativeShaderResourceDependencies();