add srt walker

CMakeLists.txt

@@ -399,7 +399,7 @@ add_executable(scalar_provenance_tests EXCLUDE_FROM_ALL
 	"${KYTY_SOURCE_DIR}/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp"
 	"${KYTY_SOURCE_DIR}/graphics/shader/recompiler/ir/passes/SrtWalker.cpp"
 )
-target_link_libraries(scalar_provenance_tests fmt::fmt)
+target_link_libraries(scalar_provenance_tests fmt::fmt xbyak::xbyak)
 target_include_directories(scalar_provenance_tests PRIVATE ${inc_headers})
 
 add_executable(page_manager_tests EXCLUDE_FROM_ALL
@@ -467,7 +467,7 @@ add_executable(resource_materialization_tests EXCLUDE_FROM_ALL
 	"${KYTY_SOURCE_DIR}/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp"
 	"${KYTY_SOURCE_DIR}/graphics/shader/recompiler/ir/passes/SrtWalker.cpp"
 )
-target_link_libraries(resource_materialization_tests fmt::fmt)
+target_link_libraries(resource_materialization_tests fmt::fmt xbyak::xbyak)
 target_include_directories(resource_materialization_tests PRIVATE ${inc_headers})
 
 add_executable(resource_tracking_tests EXCLUDE_FROM_ALL
@@ -480,7 +480,7 @@ add_executable(resource_tracking_tests EXCLUDE_FROM_ALL
 	"${KYTY_SOURCE_DIR}/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp"
 	"${KYTY_SOURCE_DIR}/graphics/shader/recompiler/ir/passes/ShaderInfoCollection.cpp"
 )
-target_link_libraries(resource_tracking_tests fmt::fmt)
+target_link_libraries(resource_tracking_tests fmt::fmt xbyak::xbyak)
 target_include_directories(resource_tracking_tests PRIVATE ${inc_headers})
 
 add_executable(audio_out2_port_tests EXCLUDE_FROM_ALL

src/graphics/shader/recompiler/ir/Program.cpp

@@ -3,6 +3,7 @@
 #include "graphics/shader/recompiler/BufferFormat.h"
 
 #include <fmt/format.h>
+#include <xbyak/xbyak.h>
 #include <map>
 #include <new>
 #include <unordered_map>
@@ -99,6 +100,9 @@ bool EquivalentValue(const ResourcePlan& program, Value left, Value right,
 
 } // namespace
 
+ResourcePlan::ResourcePlan() = default;
+ResourcePlan::ResourcePlan(ResourcePlan&&) noexcept = default;
+
 ResourcePlan::~ResourcePlan() {
 	for (auto& inst: value_storage) {
 		inst.Invalidate();

src/graphics/shader/recompiler/ir/ShaderIR.h

@@ -21,6 +21,10 @@
 #include <string_view>
 #include <vector>
 
+namespace Xbyak {
+class CodeGenerator;
+}
+
 namespace Libs::Graphics::ShaderRecompiler::IR {
 
 enum class ResourceKind {
@@ -573,12 +577,12 @@ struct ResourcePlan {
 		uint64_t           generation = 0;
 	};
 
-	ResourcePlan() = default;
+	ResourcePlan();
 	~ResourcePlan();
 
 	ResourcePlan(const ResourcePlan&)            = delete;
 	ResourcePlan& operator=(const ResourcePlan&) = delete;
-	ResourcePlan(ResourcePlan&&) noexcept         = default;
+	ResourcePlan(ResourcePlan&&) noexcept;
 	ResourcePlan& operator=(ResourcePlan&& other) noexcept;
 
 	ShaderType                    stage           = ShaderType::Unknown;
@@ -597,6 +601,7 @@ struct ResourcePlan {
 	bool                                resource_tracking_complete = false;
 	ShaderInfo                          info;
 	UniformFillPlan                     uniform_fill;
+	std::unique_ptr<Xbyak::CodeGenerator> srt_code;
 	// GPU-thread scratch for nested clean/EXEC memos, activity and material keys.
 	mutable std::deque<EvaluationContext> evaluation_contexts;
 	mutable uint32_t                       evaluation_value_count = 0;

src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp

@@ -1064,6 +1064,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 	}
 	if (capture_indirect_reads || !plan.source_reads.empty())
 		plan.resource_tracking_complete &= !program.has_address_writes;
+	CompileSrtWalker(plan);
 	return plan;
 }
 

src/graphics/shader/recompiler/ir/passes/SrtWalker.cpp

@@ -1,12 +1,17 @@
 #include "graphics/shader/recompiler/ir/passes/SrtWalker.h"
 
+#include "common/abi.h"
 #include "common/assert.h"
 #include "graphics/shader/recompiler/ir/ShaderIR.h"
 
+#include <xbyak/xbyak.h>
+
 #include <algorithm>
 #include <bit>
 #include <cmath>
 #include <cstring>
+#include <map>
+#include <unordered_map>
 #include <unordered_set>
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
@@ -359,6 +364,131 @@ private:
 
 } // namespace
 
+void CompileSrtWalker(ResourcePlan& program) {
+	if (!program.resource_tracking_complete || !program.srt_plan_complete ||
+	    program.capture_specialization_reads || !program.control_flow.empty() ||
+	    program.srt_reads.empty()) return;
+
+	using Reads = std::map<int32_t, const SrtRead*>;
+	std::map<uint32_t, Reads> roots;
+	std::unordered_map<const Inst*, Reads> pointer_uses;
+	std::unordered_map<const Inst*, int32_t> offsets;
+	for (const auto& read: program.srt_reads) {
+		const auto* inst = read.value.Resolve().TryInstruction();
+		if (inst == nullptr || inst->GetOpcode() != ValueOpcode::LoadAddressU32 ||
+		    inst->NumArgs() < 2 || !IsRawRead(program, *inst) ||
+		    read.flat_offset >= program.srt_reads.size()) return;
+		const auto offset = inst->Arg(1).Resolve();
+		if (!offset.IsImmediate() || offset.GetType() != Type::U32) return;
+		const auto& memory = program.memory_info[inst->Flags<SrtReadFlags>().index];
+		const auto relative =
+		    (static_cast<int64_t>(static_cast<int32_t>(memory.offset)) & ~int64_t {3}) +
+		    static_cast<int64_t>(offset.U32() & ~3u);
+		if (relative < INT32_MIN || relative > INT32_MAX ||
+		    !offsets.emplace(inst, static_cast<int32_t>(relative)).second) return;
+	}
+	uint32_t user_data_count = 0;
+	for (const auto& read: program.srt_reads) {
+		const auto* inst = read.value.ResolveInstruction();
+		const auto* handle = inst->Arg(0).Resolve().TryInstruction();
+		if (handle == nullptr || handle->GetOpcode() != ValueOpcode::GetAddressResource ||
+		    handle->NumArgs() != 2) return;
+		const auto* low = handle->Arg(0).Resolve().TryInstruction();
+		const auto* high = handle->Arg(1).Resolve().TryInstruction();
+		if (low == nullptr || high == nullptr) return;
+		Reads* uses = nullptr;
+		if (low->GetOpcode() == ValueOpcode::GetUserData &&
+		    high->GetOpcode() == ValueOpcode::GetUserData && low->NumArgs() == 1 &&
+		    high->NumArgs() == 1 && low->Arg(0).GetType() == Type::ScalarReg &&
+		    high->Arg(0).GetType() == Type::ScalarReg) {
+			const auto reg = RegIndex(low->Arg(0).ScalarRegister());
+			if (reg < program.user_data_base ||
+			    RegIndex(high->Arg(0).ScalarRegister()) != reg + 1u) return;
+			const auto base = reg - program.user_data_base;
+			user_data_count = std::max(user_data_count, base + 2u);
+			uses = &roots[base];
+		} else {
+			const auto lo = offsets.find(low);
+			const auto hi = offsets.find(high);
+			if (lo == offsets.end() || hi == offsets.end() ||
+			    static_cast<int64_t>(lo->second) + 4 != hi->second ||
+			    !EquivalentValue(program, low->Arg(0), high->Arg(0))) return;
+			uses = &pointer_uses[low];
+		}
+		if (!uses->emplace(offsets.at(inst), &read).second) return;
+	}
+	std::unordered_set<const Inst*> visited;
+	const auto reachable = [&](auto&& self, const Reads& reads) -> bool {
+		for (const auto& [offset, read]: reads) {
+			const auto* inst = read->value.ResolveInstruction();
+			const auto child = pointer_uses.find(inst);
+			if (child != pointer_uses.end() &&
+			    (!visited.insert(inst).second || !self(self, child->second))) return false;
+		}
+		return true;
+	};
+	for (const auto& [reg, reads]: roots) {
+		if (!reachable(reachable, reads)) return;
+	}
+	if (visited.size() != pointer_uses.size()) return;
+
+	program.srt_code = std::make_unique<Xbyak::CodeGenerator>(4096, Xbyak::AutoGrow);
+	auto& code = *program.srt_code;
+	Xbyak::Label failed, finished;
+	code.push(code.rbp);
+	code.mov(code.rbp, code.rsp);
+	code.cmp(code.r8, user_data_count);
+	code.jb(failed, Xbyak::CodeGenerator::T_NEAR);
+	code.mov(code.r9, AddressMask & ~uint64_t {3});
+	// Follow shadPS4's VisitPointer: copy one table level in offset order, then
+	// descend through its child pointers. Seed the descriptor evaluator's memo.
+	const auto visit = [&](auto&& self, int32_t pointer_offset, const Reads& reads) -> void {
+		code.push(code.rdi);
+		code.mov(code.rdi, code.qword[code.rdi + pointer_offset]);
+		code.and_(code.rdi, code.r9);
+		const auto first = reads.begin()->first;
+		const auto last = reads.rbegin()->first;
+		if (first < 0) {
+			code.mov(code.r10, -static_cast<int64_t>(first));
+			code.cmp(code.rdi, code.r10);
+			code.jb(failed, Xbyak::CodeGenerator::T_NEAR);
+		}
+		if (last > 0) {
+			code.mov(code.r10, AddressMask - last);
+			code.cmp(code.rdi, code.r10);
+			code.ja(failed, Xbyak::CodeGenerator::T_NEAR);
+		}
+		for (const auto& [offset, read]: reads) {
+			const auto index = read->value.ResolveInstruction()->EvaluationIndex(program.evaluation_value_count);
+			using Entry = ResourcePlan::EvaluationContext::Entry;
+			const auto slot = code.rdx + index * sizeof(Entry);
+			code.mov(code.r10d, code.dword[code.rdi + offset]);
+			code.mov(code.dword[code.rsi + read->flat_offset * sizeof(uint32_t)], code.r10d);
+			code.mov(code.qword[slot + offsetof(Entry, value)], code.r10);
+			code.mov(code.qword[slot + offsetof(Entry, generation)], code.rcx);
+		}
+		for (const auto& [offset, read]: reads) {
+			const auto child = pointer_uses.find(read->value.ResolveInstruction());
+			if (child == pointer_uses.end()) continue;
+			self(self, offset, child->second);
+		}
+		code.pop(code.rdi);
+	};
+	for (const auto& [reg, reads]: roots) {
+		visit(visit, reg * sizeof(uint32_t), reads);
+	}
+	code.mov(code.eax, 1);
+	code.jmp(finished, Xbyak::CodeGenerator::T_NEAR);
+	code.L(failed);
+	code.xor_(code.eax, code.eax);
+	code.L(finished);
+	code.mov(code.rsp, code.rbp);
+	code.pop(code.rbp);
+	code.ret();
+	code.readyRE();
+	EXIT_IF(Xbyak::GetError() != 0);
+}
+
 SrtWalker::SrtWalker(const ResourcePlan& program, const SrtRuntime& runtime,
                      SrtWalker* clean_evaluator, Value active_mask)
     : m_program(program), m_runtime(runtime), m_clean_evaluator(clean_evaluator),
@@ -954,6 +1084,15 @@ bool SrtWalker::EvaluateDescriptor(uint32_t source, DescriptorValue& result) {
 
 bool SrtWalker::RefreshFlatBuffer(std::vector<uint32_t>& flat) {
 	if (!m_program.srt_plan_complete) return false;
+	if (m_program.srt_code && m_runtime.read_memory == nullptr) {
+		flat.resize(m_program.srt_reads.size());
+		m_program.active_sources.clear();
+		m_context.values.resize(m_program.evaluation_value_count);
+		using Walker = bool KYTY_SYSV_ABI (*)(const uint32_t*, uint32_t*,
+		    ResourcePlan::EvaluationContext::Entry*, uint64_t, size_t);
+		return m_program.srt_code->getCode<Walker>()(m_runtime.user_data.data(), flat.data(),
+		    m_context.values.data(), m_context.generation, m_runtime.user_data.size());
+	}
 	const auto refresh = [&](uint32_t slot) {
 		if (slot >= m_program.srt_reads.size()) return false;
 		const auto& read = m_program.srt_reads[slot];

src/graphics/shader/recompiler/ir/passes/SrtWalker.h

@@ -45,6 +45,7 @@ enum class RuntimeValueType { Any, Integer, ImmutableInteger };
 
 bool ValidateRuntimeValue(const ResourcePlan& program, Value value,
                           RuntimeValueType type = RuntimeValueType::Any);
+void CompileSrtWalker(ResourcePlan& program);
 // Uses the strict reader for values that affect shader specialization.
 SrtRuntime CleanRuntime(SrtRuntime runtime);
 

tests/ResourceMaterializationTests.cpp

@@ -168,6 +168,117 @@ void TestMappedSrtUsesDirectReaderByDefault() {
         "cache rematerialization did not use the direct reader by default");
 }
 
+void TestNativeNestedSrtRefresh() {
+  using namespace Libs::Graphics::ShaderRecompiler::IR;
+  std::array<uint32_t, 3> first{0u, 0x12345678u, 0xabcdef01u};
+  std::array<uint32_t, 3> second{0u, 0x87654321u, 0x10fedcbau};
+  const auto pointer_words = [](const void *pointer) {
+    const auto address = reinterpret_cast<uint64_t>(pointer);
+    return std::array<uint32_t, 2>{static_cast<uint32_t>(address) | 3u,
+                                  static_cast<uint32_t>(address >> 32u) | 0xabcd0000u};
+  };
+  auto table = pointer_words(first.data());
+  const auto user_data = pointer_words(table.data());
+  auto program = SrtProgram(0u);
+  program.user_data_base = 4;
+  program.user_data_count = 2;
+  auto &block = *program.blocks[0];
+  const auto emit = [&](ValueOpcode opcode, std::initializer_list<Value> args) {
+    return Value(&block.AppendNewInst(opcode, args));
+  };
+  const auto low = emit(ValueOpcode::GetUserData, {Value(ScalarReg(4))});
+  const auto high = emit(ValueOpcode::GetUserData, {Value(ScalarReg(5))});
+  const auto root = emit(ValueOpcode::GetAddressResource, {low, high});
+  const auto read = [&](Value address, uint32_t offset) {
+    const auto value = emit(ValueOpcode::LoadAddressU32,
+                            {address, Value(offset), Value(0u), Value(true)});
+    value.Instruction()->SetFlags(MemoryFlags{.index = 0});
+    program.srt_reads.push_back(
+        {value, static_cast<uint32_t>(program.srt_reads.size())});
+    return value;
+  };
+  const auto pointer_low = read(root, 0u);
+  const auto pointer_high = read(root, 4u);
+  auto *leaf = program.srt_reads[0].value.Instruction();
+  leaf->Arg(0).Instruction()->SetArg(0, pointer_low);
+  leaf->Arg(0).Instruction()->SetArg(1, pointer_high);
+  leaf->SetArg(1, Value(7u));
+  read(leaf->Arg(0), 11u);
+  auto plan = ExtractResourcePlan(program);
+  Check(plan.srt_code != nullptr, "nested user-data tree did not compile an x86 walker");
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  const SrtRuntime runtime{.user_data = user_data};
+  {
+    SrtWalker walker(plan, runtime);
+    Check(walker.RefreshFlatBuffer(snapshot.flattened_srt),
+          "native SRT refresh failed");
+    const auto loaded = first[1];
+    first[1] ^= 0xffffffffu;
+    uint32_t value = 0;
+    Check(walker.Evaluate(plan.srt_reads[0].value, value) && value == loaded,
+          "descriptor evaluation reloaded a value already copied by the native walker");
+    first[1] = loaded;
+  }
+  for (const auto *data : {&first, &second}) {
+    table = pointer_words(data->data());
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+              snapshot.flattened_srt ==
+                  std::vector<uint32_t>{(*data)[1], table[0], table[1], (*data)[2]},
+          "native SRT refresh lost pointer masking, alignment, or current values");
+  }
+  Check(!MaterializeResources(plan, {.user_data = std::span(user_data).first(1)},
+                              snapshot, specialization),
+        "native SRT walk accepted truncated user data");
+  const std::array<uint32_t, 2> fake_user_data{0x1003u, 0xabcd0000u};
+  uint32_t reads = 0;
+  const SrtRuntime callback_runtime{
+      .user_data = fake_user_data,
+      .read_memory = +[](void *data, uint64_t address, std::span<uint32_t> words) {
+        ++*static_cast<uint32_t *>(data);
+        switch (address) {
+          case 0x1000u: words[0] = 0x2003u; break;
+          case 0x1004u: words[0] = 0xabcd0000u; break;
+          case 0x2004u: words[0] = 0x1234u; break;
+          case 0x2008u: words[0] = 0x5678u; break;
+          default: return false;
+        }
+        return true;
+      },
+      .userdata = &reads};
+  Check(MaterializeResources(plan, callback_runtime, snapshot, specialization) && reads == 4 &&
+            snapshot.flattened_srt ==
+                std::vector<uint32_t>{0x1234u, 0x2003u, 0xabcd0000u, 0x5678u},
+        "native-eligible SRT bypassed its callback or repeated nested pointer reads");
+  leaf->SetArg(0, root);
+  leaf->SetArg(1, Value(0u));
+  program.srt_reads.resize(1);
+  for (const int32_t offset : {-1, 7}) {
+    program.memory_info[0].offset = static_cast<uint32_t>(offset);
+    auto bounded = ExtractResourcePlan(program);
+    auto base = pointer_words(second.data() + 1);
+    Check(bounded.srt_code != nullptr &&
+              MaterializeResources(bounded, {.user_data = base}, snapshot, specialization) &&
+              snapshot.flattened_srt[0] == second[offset < 0 ? 0 : 2],
+          "native SRT applied a signed immediate offset incorrectly");
+    base = offset < 0 ? std::array<uint32_t, 2>{0u, 0u}
+                      : std::array<uint32_t, 2>{0xfffffffcu, 0xffffu};
+    Check(!MaterializeResources(bounded, {.user_data = base}, snapshot, specialization),
+          "native SRT accepted address underflow or 48-bit overflow");
+  }
+  high.Instruction()->SetArg(0, Value(ScalarReg(6)));
+  Check(ExtractResourcePlan(program).srt_code == nullptr,
+        "native SRT accepted nonadjacent user-data pointer words");
+  high.Instruction()->SetArg(0, Value(ScalarReg(5)));
+  program.memory_info[0].offset = 0;
+  program.srt_reads.push_back({pointer_low, 1});
+  program.srt_reads.push_back({pointer_high, 2});
+  leaf->SetArg(0, emit(ValueOpcode::GetAddressResource, {pointer_low, pointer_high}));
+  pointer_high.Instruction()->SetArg(1, Value(8u));
+  Check(ExtractResourcePlan(program).srt_code == nullptr,
+        "native SRT accepted nonadjacent nested pointer words");
+}
+
 void TestIntegerRuntimeValueFollowsSrtReads() {
   using namespace Libs::Graphics::ShaderRecompiler::IR;
   auto plan = SrtProgram(0x10000);
@@ -256,6 +367,8 @@ void TestSrtAliasesRetainReadPolicy() {
             plan.srt_reads[1].value.Instruction()->Flags<SrtReadFlags>().clean == 0u &&
             plan.srt_reads[2].value.Instruction()->Flags<SrtReadFlags>().clean == 1u,
         "alias normalization lost read identity or retained discarded EXEC provenance");
+  Check(plan.srt_code == nullptr,
+        "captured strict reads compiled an unused x86 walker");
   struct Reads { uint32_t strict = 0; uint32_t ordinary = 0; bool dirty = false; } reads;
   const SrtRuntime runtime{
       .read_memory = +[](void *data, uint64_t address, std::span<uint32_t> words) {
@@ -671,6 +784,7 @@ void DbgExit(int) { std::abort(); }
 
 int main() {
   TestMappedSrtUsesDirectReaderByDefault();
+  TestNativeNestedSrtRefresh();
   TestIntegerRuntimeValueFollowsSrtReads();
   TestSrtAliasesRetainReadPolicy();
   TestUniformVectorDescriptorRead();