render: set merged GS user data from shader registers

Co-authored-by: shadowbeat070 <70843316+shadowbeat070@users.noreply.github.com>

src/graphics/guest_gpu/command_processor/pm4Dispatch.cpp

@@ -139,13 +139,13 @@ constexpr auto MakeShaderDispatchTable() {
 	g_hw_sh_func[Pm4::COMPUTE_TMPRING_SIZE]            = HwShSetCsRegisters;
 	g_hw_sh_func[Pm4::COMPUTE_PGM_RSRC3]               = HwShSetCsRegisters;
 	g_hw_sh_func[Pm4::COMPUTE_PACE_ID]                 = HwShSetCsRegisters;
-	g_hw_sh_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_PS]  = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_PS]  = HwShSetRegisters;
-	g_hw_sh_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_GS]  = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_GS]  = HwShSetRegisters;
-	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_LO_GS] = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_LO_GS] = HwShSetRegisters;
-	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_HI_GS] = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_HI_GS] = HwShSetRegisters;
-	g_hw_sh_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_HS]  = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_HS]  = HwShSetRegisters;
-	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_LO_HS] = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_LO_HS] = HwShSetRegisters;
-	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_HI_HS] = HwShIgnoreRegisters;
+	g_hw_sh_func[Pm4::SPI_SHADER_USER_DATA_ADDR_HI_HS] = HwShSetRegisters;
 
 	return g_hw_sh_func;
 }

src/graphics/guest_gpu/command_processor/pm4Dispatch.h

@@ -98,7 +98,7 @@ uint32_t HwShSetCsUserSgpr(CommandProcessor&, uint32_t, uint32_t, const uint32_t
 uint32_t HwShIgnoreUserAccumulator(CommandProcessor&, uint32_t, uint32_t, const uint32_t*,
                                    uint32_t);
 uint32_t HwShSetCsRegisters(CommandProcessor&, uint32_t, uint32_t, const uint32_t*, uint32_t);
-uint32_t HwShIgnoreRegisters(CommandProcessor&, uint32_t, uint32_t, const uint32_t*, uint32_t);
+uint32_t HwShSetRegisters(CommandProcessor&, uint32_t, uint32_t, const uint32_t*, uint32_t);
 uint32_t HwUcSetPrimitiveType(CommandProcessor&, uint32_t, uint32_t, const uint32_t*, uint32_t);
 uint32_t HwUcSetIndexType(CommandProcessor&, uint32_t, uint32_t, const uint32_t*, uint32_t);
 uint32_t HwUcSetObjectId(CommandProcessor&, uint32_t, uint32_t, const uint32_t*, uint32_t);

src/graphics/guest_gpu/command_processor/pm4Handlers.cpp

@@ -980,7 +980,7 @@ static void HwShIgnoreComputeRegister([[maybe_unused]] uint32_t cmd_offset,
 static void HwShIgnoreShaderRegister([[maybe_unused]] uint32_t cmd_offset,
                                      [[maybe_unused]] uint32_t value) {}
 
-KYTY_HW_SH_PARSER(HwShIgnoreRegisters) {
+KYTY_HW_SH_PARSER(HwShSetRegisters) {
 	const auto reg_num = (cmd_id >> 16u) & 0x3fffu;
 	EXIT_NOT_IMPLEMENTED(reg_num == 0);
 
@@ -3629,12 +3629,13 @@ void GraphicsInitJmpTablesShIndirect() {
 	g_hw_sh_indirect_func[Pm4::SPI_GRAPHICS_SHADER_CONTROL_GS] = [](KYTY_HW_SH_INDIRECT_ARGS) {
 		HwShIgnoreShaderRegister(cmd_offset, value);
 	};
-	g_hw_sh_indirect_func[Pm4::SPI_SHADER_USER_DATA_ADDR_LO_GS] = [](KYTY_HW_SH_INDIRECT_ARGS) {
+	for (uint32_t offset = Pm4::SPI_SHADER_USER_DATA_ADDR_LO_GS;
-		HwShIgnoreShaderRegister(cmd_offset, value);
+	     offset <= Pm4::SPI_SHADER_USER_DATA_ADDR_HI_GS; offset++) {
-	};
+		g_hw_sh_indirect_func[offset] = [](KYTY_HW_SH_INDIRECT_ARGS) {
-	g_hw_sh_indirect_func[Pm4::SPI_SHADER_USER_DATA_ADDR_HI_GS] = [](KYTY_HW_SH_INDIRECT_ARGS) {
+			cp.GetShCtx().SetGsUserDataAddress(cmd_offset - Pm4::SPI_SHADER_USER_DATA_ADDR_LO_GS,
-		HwShIgnoreShaderRegister(cmd_offset, value);
+			                                   value);
-	};
+		};
+	}
 	g_hw_sh_indirect_func[Pm4::SPI_SHADER_PGM_CHKSUM_HS] = [](KYTY_HW_SH_INDIRECT_ARGS) {
 		HwShIgnoreShaderRegister(cmd_offset, value);
 	};

src/graphics/guest_gpu/hardwareContext.h

@@ -520,7 +520,8 @@ struct GsShaderResource2 {
 };
 
 struct GsStageRegisters {
-	uint64_t          data_addr = 0;
+	uint64_t          data_addr      = 0;
+	uint64_t          user_data_addr = 0;
 	GsShaderResource1 rsrc1;
 	GsShaderResource2 rsrc2;
 };
@@ -976,6 +977,12 @@ public:
 	void SetHsShaderResource1(const HsShaderResource1& rsrc1) { m_vs.hs_regs.rsrc1 = rsrc1; }
 	void SetHsShaderResource2(const HsShaderResource2& rsrc2) { m_vs.hs_regs.rsrc2 = rsrc2; }
 	void SetGsShaderBase(uint64_t addr) { m_vs.gs_regs.data_addr = addr; }
+	void SetGsUserDataAddress(uint32_t word, uint32_t value) {
+		const auto shift   = word * 32u;
+		auto&      address = m_vs.gs_regs.user_data_addr;
+		address = (address & ~(uint64_t {0xffffffffu} << shift)) |
+		          (static_cast<uint64_t>(value) << shift);
+	}
 	void SetGsShaderResource1(const GsShaderResource1& rsrc1) { m_vs.gs_regs.rsrc1 = rsrc1; }
 	void SetGsShaderResource2(const GsShaderResource2& rsrc2) { m_vs.gs_regs.rsrc2 = rsrc2; }
 

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

@@ -323,6 +323,7 @@ struct PipelineCache::ProgramCache {
 		if constexpr (std::is_same_v<InputInfo, ShaderVertexInputInfo>) {
 			options.user_data_base = 8;
 			if (stage == ShaderType::Mesh) {
+				options.user_data_base = 0;
 				options.wave_size      = input_info.mesh.wave_size;
 				options.scratch_dwords = input_info.mesh.scratch_size_dwords;
 			}

src/graphics/shader/shader.cpp

@@ -117,7 +117,7 @@ static ShaderParams GetShaderParams(uint64_t shader_addr, const char* label, uin
 	const auto code = std::span {reinterpret_cast<const uint32_t*>(shader_addr), code_words};
 	return {
 	    .code      = code,
-	    .user_data = user_data,
+	    .user_data = std::vector<uint32_t>(user_data.begin(), user_data.end()),
 	    .hash      = declared_hash != 0 ? declared_hash
 	                                    : XXH3_64bits(code.data(), code.size_bytes()),
 	};
@@ -864,6 +864,11 @@ ShaderParams PrepareProgram(const HW::VertexShaderInfo& regs, const HW::Context&
 	    GetShaderParams(regs.gs_regs.data_addr, "ShaderRecompiler GS",
 	                    GetDeclaredShaderHash(regs.gs_regs.data_addr), {}, back);
 	params.back_code         = back_params.code;
+	// Merged shaders receive the GS-back user-data pointer in s0:s1, followed
+	// by the ordinary user SGPRs at s8. Keep both in the runtime register snapshot.
+	params.user_data.insert(params.user_data.begin(), 8u, 0u);
+	params.user_data[0] = static_cast<uint32_t>(regs.gs_regs.user_data_addr);
+	params.user_data[1] = static_cast<uint32_t>(regs.gs_regs.user_data_addr >> 32u);
 	const uint64_t hashes[]  = {params.hash, back_params.hash};
 	params.hash              = XXH3_64bits(hashes, sizeof(hashes));
 	info                     = {};

src/graphics/shader/shaderCompiler.h

@@ -15,7 +15,7 @@ class UserConfig;
 
 struct ShaderParams {
 	std::span<const uint32_t> code;
-	std::span<const uint32_t> user_data;
+	std::vector<uint32_t>     user_data;
 	uint64_t                  hash = 0;
 	std::span<const uint32_t> back_code;
 

tests/shaderCfgTests.cpp

@@ -9054,6 +9054,82 @@ void TestFusedShaderHandoffPreservesRegisters() {
   Check(allocations == 1u, "fused shader omitted the back shader allocation");
 }
 
+void TestMergedShaderUserDataSnapshot() {
+  using namespace ShaderRecompiler;
+  const uint32_t front[] = {
+      EncodeMubuf0(0x1c), EncodeMubuf1(0, 2, 1), // consume front s[8:11]
+      EncodeSop1(0x20, 0, 6),                   // merged-stage handoff
+  };
+  const uint32_t back[] = {
+      EncodeSmem0(0x02, 8, 0), 125u << 25u, // s_load_dwordx4 s[8:11], s[0:1]
+      EncodeMubuf0(0x1c), EncodeMubuf1(0, 2, 1), EncodeSopp(0x01),
+  };
+  const std::array<uint32_t, 4> first_table = {0x12340000, 0, 64, 0x00027000};
+  const std::array<uint32_t, 4> second_table = {0x56780000, 0, 64, 0x00027000};
+  HW::VertexShaderInfo regs{};
+  regs.es_regs.data_addr = reinterpret_cast<uint64_t>(front);
+  regs.gs_regs.data_addr = reinterpret_cast<uint64_t>(back);
+  regs.gs_regs.user_data_addr = reinterpret_cast<uint64_t>(first_table.data());
+  regs.gs_regs.rsrc1.gs_vgpr_component_count = 3;
+  regs.gs_regs.rsrc2.es_vgpr_component_count = 3;
+  regs.gs_regs.rsrc2.user_sgpr = 4;
+  for (uint32_t i = 0; i < 4; i++) {
+    regs.gs_user_sgpr.value[i] = 0x10001000u + i;
+  }
+  const std::array<uint32_t, 4> front_data = {regs.gs_user_sgpr.value[0],
+      regs.gs_user_sgpr.value[1], regs.gs_user_sgpr.value[2], regs.gs_user_sgpr.value[3]};
+  ShaderMappedData mapped{};
+  mapped.code_size_bytes = sizeof(front);
+  ShaderMapUserData(regs.es_regs.data_addr, mapped);
+  mapped.code_size_bytes = sizeof(back);
+  ShaderMapUserData(regs.gs_regs.data_addr, mapped);
+  HW::Context context;
+  context.SetShaderStages(0x20);
+  context.SetMaxOutputPerSubgroup(252);
+  context.SetGsMaxVertOut(18);
+  context.SetGsOutPrimType(2);
+  HW::UserConfig user_config;
+  user_config.SetPrimitiveType(Prospero::PrimitiveType::kTriList);
+  user_config.SetGeControl({14, 14});
+  ShaderVertexInputInfo first_input{};
+  const auto first = PrepareProgram(regs, context, user_config, first_input);
+  regs.gs_regs.user_data_addr = reinterpret_cast<uint64_t>(second_table.data());
+  regs.gs_user_sgpr.value[0]++;
+  ShaderVertexInputInfo second_input{};
+  const auto second = PrepareProgram(regs, context, user_config, second_input);
+  Check(first.user_data.size() == 12 && second.user_data.size() == 12 &&
+            std::equal(front_data.begin(), front_data.end(), first.user_data.begin() + 8) &&
+            second.user_data[8] == front_data[0] + 1,
+        "merged shader parameters did not snapshot ordinary user SGPRs at s8");
+  Check(first.hash == second.hash &&
+            MakeStageStaticKey(first_input) == MakeStageStaticKey(second_input),
+        "a dynamic merged-shader user-data pointer changed shader identity");
+  CompileOptions options{};
+  options.stage = ShaderType::Mesh;
+  options.user_data_base = 0;
+  options.user_data = first.user_data;
+  options.input_info.vertex = &first_input;
+  options.back_code = first.back_code;
+  const auto translated = TranslateProgram(first.code, options);
+  const auto &program = translated.program;
+  Check(program.info.buffers.size() == 2 && program.srt_reads.size() == 4,
+        "merged shader lost front user SGPRs or the back-stage SRT load");
+  for (const auto *params : {&first, &second}) {
+    const IR::SrtRuntime runtime{.user_data = params->user_data,
+                                 .read_memory = ReadHostTestMemory};
+    IR::DescriptorValue front_descriptor, back_descriptor;
+    const auto &table = params == &first ? first_table : second_table;
+    Check(IR::EvaluateDescriptorSource(program, program.info.buffers[0].source,
+                                       runtime, front_descriptor) &&
+              IR::EvaluateDescriptorSource(program, program.info.buffers[1].source,
+                                             runtime, back_descriptor) &&
+              std::equal(params->user_data.begin() + 8, params->user_data.end(),
+                         front_descriptor.dwords.begin()) &&
+              std::equal(table.begin(), table.end(), back_descriptor.dwords.begin()),
+          "merged shader resource plan did not follow the current s0:s1 pointer and s8 data");
+  }
+}
+
 void TestMeshInputAssembly() {
   using namespace ShaderRecompiler;
   using namespace ShaderRecompiler::IR;
@@ -12406,6 +12482,7 @@ int main() {
   TestNewShaderRecompilerBranchConditionForms();
   TestNewShaderRecompilerSetpcBranch();
   TestFusedShaderHandoffPreservesRegisters();
+  TestMergedShaderUserDataSnapshot();
   TestMeshInputAssembly();
   TestNewShaderRecompilerSetpcJumpTable();
   TestNewShaderRecompilerPrunesUnreachableSetpcMetadata();